diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..2f885a478954631198a950c89a87e319cd3ce6ef --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +__pycache__/ +*.pyc +.DS_Store +loopgen-repo/ diff --git a/README.md b/README.md index 3e434c4fbd2f095d0c9b8402ee785cd809610cb3..fc0c32fd2e1ba0593d8bb2d387c71cc4a5e26f36 100644 --- a/README.md +++ b/README.md @@ -1,13 +1,14 @@ --- -title: Loopgen -emoji: 🏃 -colorFrom: red -colorTo: indigo +title: LoopGen +emoji: 🔁 +colorFrom: blue +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 +license: mit --- -Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference +Training-free loopable music generation from a text description, using MusicGen + MAGNeT. Ported from [LoopGen](https://github.com/gladia-research-group/loopgen) (Marincione et al., 2025) for use with [HARP](https://github.com/TEAMuP-dev/HARP). diff --git a/SOURCES.md b/SOURCES.md new file mode 100644 index 0000000000000000000000000000000000000000..f1c3ce005decb09a825facaf73418e70d2f1f2b3 --- /dev/null +++ b/SOURCES.md @@ -0,0 +1,4 @@ +# Sources — loopgen + +- Source repo: https://github.com/gladia-research-group/loopgen.git +- Paper: https://arxiv.org/abs/2504.04466 diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..248b0eef1b84ea185dedf96e81642960288e3e42 --- /dev/null +++ b/app.py @@ -0,0 +1,257 @@ +try: + import spaces +except ImportError: + # keep @spaces.GPU usable as a no-op; ZeroGPU requires this exact name. + 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 sys +sys.stdout.reconfigure(line_buffering=True) + +import contextlib +import gc +import os +import tempfile +import threading + +import gradio as gr +import numpy as np +import torch +from huggingface_hub import hf_hub_download +from pyharp import ModelCard, build_endpoint + +from beat_this.inference import File2Beats +from src.audiocraft.data.audio import audio_write +from src.audiocraft.models import MusicGen +from src.loopgen import LoopGen + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + +MUSICGEN_MODEL = "facebook/musicgen-medium" +MAGNET_MODEL = "facebook/magnet-medium-30secs" +MAX_CFG = 5.0 # classifier-free guidance ceiling (default: 5.0, per paper) +MIN_CFG = 1.0 # classifier-free guidance floor, annealed down to this (default: 1.0, per paper) +RESCORE_WEIGHT = 0.5 # MAGNeT/MusicGen rescoring interpolation weight (default: 0.5, per paper) +HINT_DURATION = 30.0 # seconds of MusicGen hint audio generated for beat detection (default: 30.0, per repo) +MAX_RETRIES = 5 # beat-alignment retries before falling back to a fixed length (default: 5, per repo) + +beat_checkpoint_path = None +beat_checkpoint_ready = False +beat_checkpoint_error = None + +musicgen = None +loopgen = None +file2beats = None +model_lock = threading.Lock() + + +def download_beat_checkpoint(): + """Fetch the beat-tracking checkpoint in the background so the server can + start immediately. Reuses the already-deployed teamup-tech/beat-this + Space's copy instead of re-uploading the same file here.""" + global beat_checkpoint_path, beat_checkpoint_ready, beat_checkpoint_error + try: + beat_checkpoint_path = hf_hub_download( + repo_id="teamup-tech/beat-this", filename="final0.ckpt", repo_type="space" + ) + print("Beat-tracking checkpoint downloaded.") + except Exception as e: + beat_checkpoint_error = str(e) + print(f"Beat checkpoint download error: {e}") + finally: + beat_checkpoint_ready = True + + +threading.Thread(target=download_beat_checkpoint, daemon=True).start() + + +model_card = ModelCard( + name="LoopGen", + description="Generates a seamlessly loopable music clip from a text description, training-free, using MusicGen + MAGNeT.", + author="Davide Marincione, Giorgio Strano, Donato Crisostomi, Roberto Ribuoli, Emanuele Rodolà", + tags=["music generation", "loop", "text-to-music"], +) + + +def _write_temp_wav(wav, sample_rate): + """Write a generated waveform to a fresh temp file and return its path. + Adapted from pipelines.py, which reused hardcoded local paths like + f"musicgen_base_{name}_{seed}.wav" across calls (reused via + os.path.exists) -- only valid for one person rerunning experiments + locally, since concurrent Space users could collide on the same + seed/name and clobber or read each other's files.""" + fd, stem = tempfile.mkstemp() + os.close(fd) + os.remove(stem) # audio_write creates its own file at stem + extension + path = audio_write(stem, wav, sample_rate, loudness_compressor=True, + strategy="loudness", loudness_headroom_db=16) + return str(path) + + +def _beat_aligned_length(path, min_length, max_length, unit_size=0.02, length_retries=2): + """Find a loop length, in 20ms units, aligned to a whole number of bars, + within [min_length, max_length] seconds. Adapted from pipelines.py's + get_beat_perfect_hint, unchanged except for taking an already-constructed + file2beats (construction moved to process_fn, see its docstring).""" + min_units = min_length // unit_size + max_units = max_length // unit_size + + beats, downbeats = file2beats(path) + + diffs = [] + prev = None + beats_per_bar = [] + beats_per_this_bar = 0 + downbeat_idx = 0 + for beat in beats: + while downbeat_idx < len(downbeats) and beat > downbeats[downbeat_idx]: + beats_per_bar.append(beats_per_this_bar) + beats_per_this_bar = 0 + downbeat_idx += 1 + beats_per_this_bar += 1 + if prev is not None: + diffs.append(beat - prev) + prev = beat + + if len(beats_per_bar) < 1: + return None + + diffs = sorted(diffs) + beats_per_bar = sorted(beats_per_bar) + + beat_length = diffs[len(diffs) // 2] // unit_size + 1 # quantize the median to unit_size + beats_per_bar = beats_per_bar[len(beats_per_bar) // 2] # median beats-per-bar + bar_length = beats_per_bar * beat_length + length = bar_length * 12 + if length <= 0: + return None + + loops_done = 0 + while length < min_units or length > max_units: + if length < min_units: + length *= 2 + else: + length /= 2 + loops_done += 1 + + if loops_done > length_retries: + return None + return length + + +@spaces.GPU(duration=300) +@torch.inference_mode() +def process_fn(description: str, min_length: float, max_length: float, seed: int) -> str: + """Generates a seamlessly loopable music clip: a MusicGen hint is used to + find a beat-aligned loop length (retrying up to MAX_RETRIES times), then + MAGNeT generates the final loop as a continuation of that hint. Adapted + from pipelines.py's loopgen_loop. + + Model construction happens here, inside the GPU call, rather than in the + background-download thread above: LoopGen/MusicGen.get_pretrained() + bakes device="cuda" in immediately, and ZeroGPU only allows CUDA calls + made inside an @spaces.GPU-decorated call, not from a background thread.""" + global musicgen, loopgen, file2beats + + if not beat_checkpoint_ready: + raise gr.Error("Beat-tracking checkpoint is still downloading, please try again shortly.") + if beat_checkpoint_error is not None: + raise gr.Error(f"Beat-tracking checkpoint failed to download: {beat_checkpoint_error}") + + with model_lock: + if musicgen is None: + musicgen = MusicGen.get_pretrained(MUSICGEN_MODEL, device=DEVICE) + if loopgen is None: + loopgen = LoopGen.get_pretrained(MAGNET_MODEL, device=DEVICE) + if file2beats is None: + file2beats = File2Beats(checkpoint_path=beat_checkpoint_path, device=DEVICE, dbn=False) + + torch.manual_seed(seed) + np.random.seed(seed) + + unit_length = None + hint = None + sr = None + tries = 0 + while unit_length is None: + gc.collect() + if DEVICE == "cuda": + torch.cuda.empty_cache() + if tries >= MAX_RETRIES: + unit_length = 350 * 3 + break + del hint + tries += 1 + musicgen.set_generation_params(duration=HINT_DURATION) + hint = musicgen.generate([description])[0].to(device="cpu", dtype=torch.float32) + sr = musicgen.sample_rate + hint_path = _write_temp_wav(hint, sr) + unit_length = _beat_aligned_length(hint_path, min_length, max_length) + os.remove(hint_path) + + gc.collect() + if DEVICE == "cuda": + torch.cuda.empty_cache() + + left_hint = [hint[..., :int(0.02 * sr) * int(unit_length / 2)]] + valid_tokens = [int(unit_length)] + + setup = { + "span_arrangement": "stride1", + "use_sampling": True, + "top_k": 0, + "top_p": 0.9, + "temperature": 3.0, + "max_cfg_coef": MAX_CFG, + "min_cfg_coef": MIN_CFG, + "decoding_steps": [100, 50, 10, 10], + "rescorer": musicgen.lm, + "rescore_weights": RESCORE_WEIGHT, + "offset": False, + } + loopgen.set_generation_params(**setup) + + autocast_ctx = (torch.autocast(device_type="cuda", dtype=torch.float16) + if DEVICE == "cuda" else contextlib.nullcontext()) + with autocast_ctx: + results = loopgen.generate_continuation( + text_prompt=description, negative_text_prompt=None, + left_hint=left_hint, left_hint_sr=sr, valid_tokens=valid_tokens, + ) + result = results[0].to(device="cpu", dtype=torch.float32) + + return _write_temp_wav(result, loopgen.sample_rate) + + +with gr.Blocks() as demo: + input_components = [ + gr.Textbox(label="Description", + info="Text description of the music to generate."), + gr.Slider(minimum=5, maximum=30, step=1, value=15, label="Min Loop Length (s)", + info="Shortest acceptable loop length in seconds (default: 15, per repo)."), + gr.Slider(minimum=5, maximum=30, step=1, value=30, label="Max Loop Length (s)", + info="Longest acceptable loop length in seconds (default: 30, per repo)."), + gr.Number(value=42, precision=0, label="Seed", + info="Random seed for generation (default: 42, per repo)."), + ] + output_components = [ + gr.Audio(type="filepath", label="Output Audio").set_info("Generated loopable music clip."), + ] + + build_endpoint( + model_card=model_card, + input_components=input_components, + output_components=output_components, + process_fn=process_fn, + ) + +if __name__ == "__main__": + demo.queue().launch(pwa=True) diff --git a/beat_this/__init__.py b/beat_this/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/beat_this/inference.py b/beat_this/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..53c198dde2f88dc74c4cc6208bcf8d6f64a27f24 --- /dev/null +++ b/beat_this/inference.py @@ -0,0 +1,315 @@ +import inspect + +import numpy as np +import soxr +import torch +import torch.nn.functional as F + +from beat_this.model.beat_tracker import BeatThis +from beat_this.model.postprocessor import Postprocessor +from beat_this.preprocessing import LogMelSpect, load_audio +from beat_this.utils import replace_state_dict_key, save_beat_tsv + +CHECKPOINT_URL = "https://cloud.cp.jku.at/public.php/dav/files/7ik4RrBKTS273gp" + + +def load_checkpoint(checkpoint_path: str, device: str | torch.device = "cpu") -> dict: + """ + Load a BeatThis checkpoint as a dictionary. + + Args: + checkpoint_path (str, optional): The path to the checkpoint. Can be a local path, a URL, or a shortname. + device (torch.device or str): The device to load the model on. + + Returns: + dict: The loaded checkpoint dictionary. + """ + try: + # try interpreting as local file name + weights_only = {"weights_only": True} if torch.__version__ >= "2" else {} + return torch.load(checkpoint_path, map_location=device, **weights_only) + except FileNotFoundError: + try: + if not ( + str(checkpoint_path).startswith("https://") + or str(checkpoint_path).startswith("http://") + ): + # interpret it as a name of one of our checkpoints + checkpoint_url = f"{CHECKPOINT_URL}/{checkpoint_path}.ckpt" + file_name = f"beat_this-{checkpoint_path}.ckpt" + else: + # try interpreting as a URL + checkpoint_url = checkpoint_path + file_name = None + return torch.hub.load_state_dict_from_url( + checkpoint_url, + file_name=file_name, + map_location=device, + ) + except Exception: + raise ValueError( + "Could not load the checkpoint given the provided name", + checkpoint_path, + ) + + +def load_model( + checkpoint_path: str | None = "final0", device: str | torch.device = "cpu" +) -> BeatThis: + """ + Load a BeatThis model from a checkpoint. + + Args: + checkpoint_path (str, optional): The path to the checkpoint. Can be a local path, a URL, or a shortname. + device (torch.device or str): The device to load the model on. + + Returns: + BeatThis: The loaded model. + """ + if checkpoint_path is not None: + checkpoint = load_checkpoint(checkpoint_path, device) + # Retrieve the model hyperparameters as it could be the small model + hparams = checkpoint["hyper_parameters"] + # Filter only those hyperparameters that apply to the model itself + hparams = { + k: v + for k, v in hparams.items() + if k in set(inspect.signature(BeatThis).parameters) + } + # Create the uninitialized model + model = BeatThis(**hparams) + # The PLBeatThis (LightningModule) state_dict contains the BeatThis + # state_dict under the "model." prefix; remove the prefix to load it + state_dict = replace_state_dict_key(checkpoint["state_dict"], "model.", "") + model.load_state_dict(state_dict) + else: + model = BeatThis() + return model.to(device).eval() + + +def zeropad(spect: torch.Tensor, left: int = 0, right: int = 0): + """ + Pads a tensor spectrogram matrix of shape (time x bins) with `left` frames in the beginning and `right` frames in the end. + """ + if left == 0 and right == 0: + return spect + else: + return F.pad(spect, (0, 0, left, right), "constant", 0) + + +def split_piece( + spect: torch.Tensor, + chunk_size: int, + border_size: int = 6, + avoid_short_end: bool = True, +): + """ + Split a tensor spectrogram matrix of shape (time x bins) into time chunks of `chunk_size` and return the chunks and starting positions. + The `border_size` is the number of frames assumed to be discarded in the predictions on either side (since the model was not trained on the input edges due to the max-pool in the loss). + To cater for this, the first and last chunk are padded by `border_size` on the beginning and end, respectively, and consecutive chunks overlap by `border_size`. + If `avoid_short_end` is true, the last chunk start is shifted left to ends at the end of the piece, therefore the last chunk can potentially overlap with previous chunks more than border_size, otherwise it will be a shorter segment. + If the piece is shorter than `chunk_size`, avoid_short_end is ignored and the piece is returned as a single shorter chunk. + + Args: + spect (torch.Tensor): The input spectrogram tensor of shape (time x bins). + chunk_size (int): The size of the chunks to produce. + border_size (int, optional): The size of the border to overlap between chunks. Defaults to 6. + avoid_short_end (bool, optional): If True, the last chunk is shifted left to end at the end of the piece. Defaults to True. + """ + # generate the start and end indices + starts = np.arange( + -border_size, len(spect) - border_size, chunk_size - 2 * border_size + ) + if avoid_short_end and len(spect) > chunk_size - 2 * border_size: + # if we avoid short ends, move the last index to the end of the piece - (chunk_size - border_size) + starts[-1] = len(spect) - (chunk_size - border_size) + # generate the chunks + chunks = [ + zeropad( + spect[max(start, 0) : min(start + chunk_size, len(spect))], + left=max(0, -start), + right=max(0, min(border_size, start + chunk_size - len(spect))), + ) + for start in starts + ] + return chunks, starts + + +def aggregate_prediction( + pred_chunks: list, + starts: list, + full_size: int, + chunk_size: int, + border_size: int, + overlap_mode: str, + device: str | torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Aggregates the predictions for the whole piece based on the given prediction chunks. + + Args: + pred_chunks (list): List of prediction chunks, where each chunk is a dictionary containing 'beat' and 'downbeat' predictions. + starts (list): List of start positions for each prediction chunk. + full_size (int): Size of the full piece. + chunk_size (int): Size of each prediction chunk. + border_size (int): Size of the border to be discarded from each prediction chunk. + overlap_mode (str): Mode for handling overlapping predictions. Can be 'keep_first' or 'keep_last'. + device (torch.device): Device to be used for the predictions. + + Returns: + tuple: A tuple containing the aggregated beat predictions and downbeat predictions as torch tensors for the whole piece. + """ + if border_size > 0: + # cut the predictions to discard the border + pred_chunks = [ + { + "beat": pchunk["beat"][border_size:-border_size], + "downbeat": pchunk["downbeat"][border_size:-border_size], + } + for pchunk in pred_chunks + ] + # aggregate the predictions for the whole piece + piece_prediction_beat = torch.full((full_size,), -1000.0, device=device) + piece_prediction_downbeat = torch.full((full_size,), -1000.0, device=device) + if overlap_mode == "keep_first": + # process in reverse order, so predictions of earlier excerpts overwrite later ones + pred_chunks = reversed(list(pred_chunks)) + starts = reversed(list(starts)) + for start, pchunk in zip(starts, pred_chunks): + piece_prediction_beat[ + start + border_size : start + chunk_size - border_size + ] = pchunk["beat"] + piece_prediction_downbeat[ + start + border_size : start + chunk_size - border_size + ] = pchunk["downbeat"] + return piece_prediction_beat, piece_prediction_downbeat + + +def split_predict_aggregate( + spect: torch.Tensor, + chunk_size: int, + border_size: int, + overlap_mode: str, + model: torch.nn.Module, +) -> dict: + """ + Function for pieces that are longer than the training length of the model. + Split the input piece into chunks, run the model on them, and aggregate the predictions. + The spect is supposed to be a torch tensor of shape (time x bins), i.e., unbatched, and the output is also unbatched. + + Args: + spect (torch.Tensor): the input piece + chunk_size (int): the length of the chunks + border_size (int): the size of the border that is discarded from the predictions + overlap_mode (str): how to handle overlaps between chunks + model (torch.nn.Module): the model to run + + Returns: + dict: the model framewise predictions for the hole piece as a dictionary containing 'beat' and 'downbeat' predictions. + """ + # split the piece into chunks + chunks, starts = split_piece( + spect, chunk_size, border_size=border_size, avoid_short_end=True + ) + # run the model + pred_chunks = [model(chunk.unsqueeze(0)) for chunk in chunks] + # remove the extra dimension in beat and downbeat prediction due to batch size 1 + pred_chunks = [ + {"beat": p["beat"][0], "downbeat": p["downbeat"][0]} for p in pred_chunks + ] + piece_prediction_beat, piece_prediction_downbeat = aggregate_prediction( + pred_chunks, + starts, + spect.shape[0], + chunk_size, + border_size, + overlap_mode, + spect.device, + ) + # save it to model_prediction + return {"beat": piece_prediction_beat, "downbeat": piece_prediction_downbeat} + + +class Spect2Frames: + """ + Class for extracting framewise beat and downbeat predictions (logits) from a spectrogram. + """ + + def __init__(self, checkpoint_path="final0", device="cpu", float16=False): + super().__init__() + self.device = torch.device(device) + self.float16 = float16 + self.model = load_model(checkpoint_path, self.device) + + def spect2frames(self, spect): + with torch.inference_mode(): + with torch.autocast(enabled=self.float16, device_type=self.device.type): + model_prediction = split_predict_aggregate( + spect=spect, + chunk_size=1500, + overlap_mode="keep_first", + border_size=6, + model=self.model, + ) + return model_prediction["beat"].float(), model_prediction["downbeat"].float() + + def __call__(self, spect): + return self.spect2frames(spect) + + +class Audio2Frames(Spect2Frames): + """ + Class for extracting framewise beat and downbeat predictions (logits) from an audio tensor. + """ + + def __init__(self, checkpoint_path="final0", device="cpu", float16=False): + super().__init__(checkpoint_path, device, float16) + self.spect = LogMelSpect(device=self.device) + + def signal2spect(self, signal, sr): + if signal.ndim == 2: + signal = signal.mean(1) + elif signal.ndim != 1: + raise ValueError(f"Expected 1D or 2D signal, got shape {signal.shape}") + if sr != 22050: + signal = soxr.resample(signal, in_rate=sr, out_rate=22050) + signal = torch.tensor(signal, dtype=torch.float32, device=self.device) + return self.spect(signal) + + def __call__(self, signal, sr): + spect = self.signal2spect(signal, sr) + return self.spect2frames(spect) + + +class Audio2Beats(Audio2Frames): + """ + Class for extracting beat and downbeat positions (in seconds) from an audio tensor. + + Args: + checkpoint_path (str): Path to the model checkpoint file. It can be a local path, a URL, or a key from the CHECKPOINT_URL dictionary. Default is "final0", which will load the model trained on all data except GTZAN with seed 0. + device (str): Device to use for inference. Default is "cpu". + float16 (bool): Whether to use half precision floating point arithmetic. Default is False. + dbn (bool): Whether to use the madmom DBN for post-processing. Default is False. + """ + + def __init__( + self, checkpoint_path="final0", device="cpu", float16=False, dbn=False + ): + super().__init__(checkpoint_path, device, float16) + self.frames2beats = Postprocessor(type="dbn" if dbn else "minimal") + + def __call__(self, signal, sr): + beat_logits, downbeat_logits = super().__call__(signal, sr) + return self.frames2beats(beat_logits, downbeat_logits) + + +class File2Beats(Audio2Beats): + def __call__(self, audio_path): + signal, sr = load_audio(audio_path) + return super().__call__(signal, sr) + + +class File2File(File2Beats): + def __call__(self, audio_path, output_path): + downbeats, beats = super().__call__(audio_path) + save_beat_tsv(downbeats, beats, output_path) diff --git a/beat_this/model/__init__.py b/beat_this/model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/beat_this/model/beat_tracker.py b/beat_this/model/beat_tracker.py new file mode 100644 index 0000000000000000000000000000000000000000..28fc63031f78fd474f425ffd7803e7de54936389 --- /dev/null +++ b/beat_this/model/beat_tracker.py @@ -0,0 +1,346 @@ +""" +Model definitions for the Beat This! beat tracker. +""" + +import contextlib +from collections import OrderedDict + +import torch +from einops import rearrange +from einops.layers.torch import Rearrange +from rotary_embedding_torch import RotaryEmbedding +from torch import nn + +from beat_this.model import roformer +from beat_this.utils import replace_state_dict_key + + +class BeatThis(nn.Module): + """ + A neural network model for beat tracking. It is composed of three main components: + - a frontend that processes the input spectrogram, + - a series of transformer blocks that process the output of the frontend, + - a head that produces the final beat and downbeat predictions. + + Args: + spect_dim (int): The dimension of the input spectrogram (default: 128). + transformer_dim (int): The dimension of the main transformer blocks (default: 512). + ff_mult (int): The multiplier for the feed-forward dimension in the transformer blocks (default: 4). + n_layers (int): The number of transformer blocks (default: 6). + head_dim (int): The dimension of each attention head for the partial transformers in the frontend and the transformer blocks (default: 32). + stem_dim (int): The out dimension of the stem convolutional layer (default: 32). + dropout (dict): A dictionary specifying the dropout rates for different parts of the model + (default: {"frontend": 0.1, "transformer": 0.2}). + sum_head (bool): Whether to use a SumHead for the final predictions (default: True) or plain independent projections. + partial_transformers (bool): Whether to include partial frequency- and time-transformers in the frontend (default: True) + """ + + def __init__( + self, + spect_dim: int = 128, + transformer_dim: int = 512, + ff_mult: int = 4, + n_layers: int = 6, + head_dim: int = 32, + stem_dim: int = 32, + dropout: dict = {"frontend": 0.1, "transformer": 0.2}, + sum_head: bool = True, + partial_transformers: bool = True, + ): + super().__init__() + # shared rotary embedding for frontend blocks and transformer blocks + rotary_embed = RotaryEmbedding(head_dim) + + # create the frontend + # - stem + stem = self.make_stem(spect_dim, stem_dim) + spect_dim //= 4 # frequencies were convolved with stride 4 + # - three frontend blocks + frontend_blocks = [] + dim = stem_dim + for _ in range(3): + frontend_blocks.append( + self.make_frontend_block( + dim, + dim * 2, + partial_transformers, + head_dim, + rotary_embed, + dropout["frontend"], + ) + ) + dim *= 2 + spect_dim //= 2 # frequencies were convolved with stride 2 + frontend_blocks = nn.Sequential(*frontend_blocks) + # - linear projection to transformer dimensionality + concat = Rearrange("b c f t -> b t (c f)") + linear = nn.Linear(dim * spect_dim, transformer_dim) + self.frontend = nn.Sequential( + OrderedDict(stem=stem, blocks=frontend_blocks, concat=concat, linear=linear) + ) + + # create the transformer blocks + assert ( + transformer_dim % head_dim == 0 + ), "transformer_dim must be divisible by head_dim" + n_heads = transformer_dim // head_dim + self.transformer_blocks = roformer.Transformer( + dim=transformer_dim, + depth=n_layers, + heads=n_heads, + attn_dropout=dropout["transformer"], + ff_dropout=dropout["transformer"], + rotary_embed=rotary_embed, + ff_mult=ff_mult, + dim_head=head_dim, + norm_output=True, + ) + + # create the output heads + if sum_head: + self.task_heads = SumHead(transformer_dim) + else: + self.task_heads = Head(transformer_dim) + + # init all weights + self.apply(self._init_weights) + + @staticmethod + def make_stem(spect_dim: int, stem_dim: int) -> nn.Module: + return nn.Sequential( + OrderedDict( + rearrange_tf=Rearrange("b t f -> b f t"), + bn1d=nn.BatchNorm1d(spect_dim), + add_channel=Rearrange("b f t -> b 1 f t"), + conv2d=nn.Conv2d( + in_channels=1, + out_channels=stem_dim, + kernel_size=(4, 3), + stride=(4, 1), + padding=(0, 1), + bias=False, + ), + bn2d=nn.BatchNorm2d(stem_dim), + activation=nn.GELU(), + ) + ) + + @staticmethod + def make_frontend_block( + in_dim: int, + out_dim: int, + partial_transformers: bool = True, + head_dim: int | None = 32, + rotary_embed: RotaryEmbedding | None = None, + dropout: float = 0.1, + ) -> nn.Module: + if partial_transformers and (head_dim is None or rotary_embed is None): + raise ValueError( + "Must specify head_dim and rotary_embed for using partial_transformers" + ) + return nn.Sequential( + OrderedDict( + partial=( + PartialFTTransformer( + dim=in_dim, + dim_head=head_dim, + n_head=in_dim // head_dim, + rotary_embed=rotary_embed, + dropout=dropout, + ) + if partial_transformers + else nn.Identity() + ), + # conv block + conv2d=nn.Conv2d( + in_channels=in_dim, + out_channels=out_dim, + kernel_size=(2, 3), + stride=(2, 1), + padding=(0, 1), + bias=False, + ), + # out_channels : 64, 128, 256 + # freqs : 16, 8, 4 (due to the stride=2) + norm=nn.BatchNorm2d(out_dim), + activation=nn.GELU(), + ) + ) + + @staticmethod + def _init_weights(module: nn.Module): + if isinstance(module, (nn.Linear, nn.Conv1d)): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + if module.bias is not None: + torch.nn.init.zeros_(module.bias) + elif isinstance(module, nn.Conv2d): + torch.nn.init.kaiming_normal_( + module.weight, mode="fan_out", nonlinearity="relu" + ) + if module.bias is not None: + torch.nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + if module.padding_idx is not None: + with torch.no_grad(): + module.weight[module.padding_idx].fill_(0) + + def forward(self, x): + x = self.frontend(x) + x = self.transformer_blocks(x) + x = self.task_heads(x) + return x + + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + # remove _orig_mod prefixes for compiled models + state_dict = replace_state_dict_key(state_dict, "_orig_mod.", "") + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + + def state_dict(self, *args, **kwargs): + state_dict = super().state_dict(*args, **kwargs) + # remove _orig_mod prefixes for compiled models + state_dict = replace_state_dict_key(state_dict, "_orig_mod.", "") + return state_dict + + +class PartialRoformer(nn.Module): + """ + Takes a (batch, channels, freqs, time) input, applies self-attention and + a feed-forward block either only across frequencies or only across time. + Returns a tensor of the same shape as the input. + """ + + def __init__( + self, + dim: int, + dim_head: int, + n_head: int, + direction: str, + rotary_embed: RotaryEmbedding, + dropout: float, + ): + super().__init__() + + assert dim % dim_head == 0, "dim must be divisible by dim_head" + assert dim // dim_head == n_head, "n_head must be equal to dim // dim_head" + self.direction = direction[0].lower() + if self.direction not in "ft": + raise ValueError(f"direction must be F or T, got {direction}") + self.attn = roformer.Attention( + dim, + heads=n_head, + dim_head=dim_head, + dropout=dropout, + rotary_embed=rotary_embed, + ) + self.ff = roformer.FeedForward(dim, dropout=dropout) + + def forward(self, x): + b = len(x) + if self.direction == "f": + pattern = "(b t) f c" + elif self.direction == "t": + pattern = "(b f) t c" + x = rearrange(x, f"b c f t -> {pattern}") + x = x + self.attn(x) + x = x + self.ff(x) + x = rearrange(x, f"{pattern} -> b c f t", b=b) + return x + + +class PartialFTTransformer(nn.Module): + """ + Takes a (batch, channels, freqs, time) input, applies self-attention and + a feed-forward block once across frequencies and once across time. Same + as applying two PartialRoformer() in sequence, but encapsulated in a single + module. Returns a tensor of the same shape as the input. + """ + + def __init__( + self, + dim: int, + dim_head: int, + n_head: int, + rotary_embed: RotaryEmbedding, + dropout: float, + ): + super().__init__() + + assert dim % dim_head == 0, "dim must be divisible by dim_head" + assert dim // dim_head == n_head, "n_head must be equal to dim // dim_head" + # frequency directed partial transformer + self.attnF = roformer.Attention( + dim, + heads=n_head, + dim_head=dim_head, + dropout=dropout, + rotary_embed=rotary_embed, + ) + self.ffF = roformer.FeedForward(dim, dropout=dropout) + # time directed partial transformer + self.attnT = roformer.Attention( + dim, + heads=n_head, + dim_head=dim_head, + dropout=dropout, + rotary_embed=rotary_embed, + ) + self.ffT = roformer.FeedForward(dim, dropout=dropout) + + def forward(self, x): + b = len(x) + # frequency directed partial transformer + x = rearrange(x, "b c f t -> (b t) f c") + x = x + self.attnF(x) + x = x + self.ffF(x) + # time directed partial transformer + x = rearrange(x, "(b t) f c ->(b f) t c", b=b) + x = x + self.attnT(x) + x = x + self.ffT(x) + x = rearrange(x, "(b f) t c -> b c f t", b=b) + return x + + +class SumHead(nn.Module): + """ + A PyTorch module that produces the final beat and downbeat prediction logits. + The beats are a sum of all beats and all downbeats predictions, to reduce the prediction + of downbeats which are not beats. + """ + + def __init__(self, input_dim): + super().__init__() + self.beat_downbeat_lin = nn.Linear(input_dim, 2) + + def forward(self, x): + beat_downbeat = self.beat_downbeat_lin(x) + # separate beat from downbeat + beat, downbeat = rearrange(beat_downbeat, "b t c -> c b t", c=2) + # aggregate beats and downbeats prediction + # autocast to float16 disabled to avoid numerical issues causing NaNs + if hasattr( + torch.amp, "is_autocast_available" + ) and not torch.amp.is_autocast_available(beat.device.type): + # but do not try disabling if the device does not support autocast + disable_autocast = contextlib.nullcontext() + else: + disable_autocast = torch.autocast(beat.device.type, enabled=False) + with disable_autocast: + beat = beat.float() + downbeat.float() + return {"beat": beat, "downbeat": downbeat} + + +class Head(nn.Module): + """ + A PyToch module that produces the final beat and downbeat prediction logits with independent linear layers outputs. + """ + + def __init__(self, input_dim): + super().__init__() + self.beat_downbeat_lin = nn.Linear(input_dim, 2) + + def forward(self, x): + beat_downbeat = self.beat_downbeat_lin(x) + # separate beat from downbeat + beat, downbeat = rearrange(beat_downbeat, "b t c -> c b t", c=2) + return {"beat": beat, "downbeat": downbeat} diff --git a/beat_this/model/postprocessor.py b/beat_this/model/postprocessor.py new file mode 100644 index 0000000000000000000000000000000000000000..45a90a6bcf831f803b8228c8951c9c3c2e90d97a --- /dev/null +++ b/beat_this/model/postprocessor.py @@ -0,0 +1,197 @@ +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +import torch +import torch.nn.functional as F +from einops import rearrange + + +class Postprocessor: + """Postprocessor for the beat and downbeat predictions of the model. + The postprocessor takes the (framewise) model predictions (beat and downbeats) and the padding mask, + and returns the postprocessed beat and downbeat as list of times in seconds. + The beats and downbeats can be 1D arrays (for only 1 piece) or 2D arrays, if a batch of pieces is considered. + The output dimensionality is the same as the input dimensionality. + Two types of postprocessing are implemented: + - minimal: a simple postprocessing that takes the maximum of the framewise predictions, + and removes adjacent peaks. + - dbn: a postprocessing based on the Dynamic Bayesian Network proposed by Böck et al. + Args: + type (str): the type of postprocessing to apply. Either "minimal" or "dbn". Default is "minimal". + fps (int): the frames per second of the model framewise predictions. Default is 50. + """ + + def __init__(self, type: str = "minimal", fps: int = 50): + assert type in ["minimal", "dbn"] + self.type = type + self.fps = fps + if type == "dbn": + from madmom.features.downbeats import DBNDownBeatTrackingProcessor + + self.dbn = DBNDownBeatTrackingProcessor( + beats_per_bar=[3, 4], + min_bpm=55.0, + max_bpm=215.0, + fps=self.fps, + transition_lambda=100, + ) + + def __call__( + self, + beat: torch.Tensor, + downbeat: torch.Tensor, + padding_mask: torch.Tensor | None = None, + ) -> tuple[np.ndarray, np.ndarray]: + """ + Apply postprocessing to the input beat and downbeat tensors. Works with batched and unbatched inputs. + The output is a list of times in seconds, or a list of lists of times in seconds, if the input is batched. + + Args: + beat (torch.Tensor): The input beat tensor. + downbeat (torch.Tensor): The input downbeat tensor. + padding_mask (torch.Tensor, optional): The padding mask tensor. Defaults to None. + + Returns: + torch.Tensor: The postprocessed beat tensor. + torch.Tensor: The postprocessed downbeat tensor. + """ + batched = False if beat.ndim == 1 else True + if padding_mask is None: + padding_mask = torch.ones_like(beat, dtype=torch.bool) + + # if beat and downbeat are 1D tensors, add a batch dimension + if not batched: + beat = beat.unsqueeze(0) + downbeat = downbeat.unsqueeze(0) + padding_mask = padding_mask.unsqueeze(0) + + if self.type == "minimal": + postp_beat, postp_downbeat = self.postp_minimal( + beat, downbeat, padding_mask + ) + elif self.type == "dbn": + postp_beat, postp_downbeat = self.postp_dbn(beat, downbeat, padding_mask) + else: + raise ValueError("Invalid postprocessing type") + + # remove the batch dimension if it was added + if not batched: + postp_beat = postp_beat[0] + postp_downbeat = postp_downbeat[0] + + # update the model prediction dict + return postp_beat, postp_downbeat + + def postp_minimal(self, beat, downbeat, padding_mask): + # concatenate beat and downbeat in the same tensor of shape (B, T, 2) + packed_pred = rearrange( + [beat, downbeat], "c b t -> b t c", b=beat.shape[0], t=beat.shape[1], c=2 + ) + # set padded elements to -1000 (= probability zero even in float64) so they don't influence the maxpool + pred_logits = packed_pred.masked_fill(~padding_mask.unsqueeze(-1), -1000) + # reshape to (2*B, T) to apply max pooling + pred_logits = rearrange(pred_logits, "b t c -> (c b) t") + # pick maxima within +/- 70ms + pred_peaks = pred_logits.masked_fill( + pred_logits != F.max_pool1d(pred_logits, 7, 1, 3), -1000 + ) + # keep maxima with over 0.5 probability (logit > 0) + pred_peaks = pred_peaks > 0 + # rearrange back to two tensors of shape (B, T) + beat_peaks, downbeat_peaks = rearrange( + pred_peaks, "(c b) t -> c b t", b=beat.shape[0], t=beat.shape[1], c=2 + ) + # run the piecewise operations + with ThreadPoolExecutor() as executor: + postp_beat, postp_downbeat = zip( + *executor.map( + self._postp_minimal_item, beat_peaks, downbeat_peaks, padding_mask + ) + ) + return postp_beat, postp_downbeat + + def _postp_minimal_item(self, padded_beat_peaks, padded_downbeat_peaks, mask): + """Function to compute the operations that must be computed piece by piece, and cannot be done in batch.""" + # unpad the predictions by truncating the padding positions + beat_peaks = padded_beat_peaks[mask] + downbeat_peaks = padded_downbeat_peaks[mask] + # pass from a boolean array to a list of times in frames. + beat_frame = torch.nonzero(beat_peaks).cpu().numpy()[:, 0] + downbeat_frame = torch.nonzero(downbeat_peaks).cpu().numpy()[:, 0] + # remove adjacent peaks + beat_frame = deduplicate_peaks(beat_frame, width=1) + downbeat_frame = deduplicate_peaks(downbeat_frame, width=1) + # convert from frame to seconds + beat_time = beat_frame / self.fps + downbeat_time = downbeat_frame / self.fps + # move the downbeat to the nearest beat + if ( + len(beat_time) > 0 + ): # skip if there are no beats, like in the first training steps + for i, d_time in enumerate(downbeat_time): + beat_idx = np.argmin(np.abs(beat_time - d_time)) + downbeat_time[i] = beat_time[beat_idx] + # remove duplicate downbeat times (if some db were moved to the same position) + downbeat_time = np.unique(downbeat_time) + return beat_time, downbeat_time + + def postp_dbn(self, beat, downbeat, padding_mask): + beat_prob = beat.double().sigmoid() + downbeat_prob = downbeat.double().sigmoid() + # limit lower and upper bound, since 0 and 1 create problems in the DBN + epsilon = 1e-5 + beat_prob = beat_prob * (1 - epsilon) + epsilon / 2 + downbeat_prob = downbeat_prob * (1 - epsilon) + epsilon / 2 + with ThreadPoolExecutor() as executor: + postp_beat, postp_downbeat = zip( + *executor.map( + self._postp_dbn_item, beat_prob, downbeat_prob, padding_mask + ) + ) + return postp_beat, postp_downbeat + + def _postp_dbn_item(self, padded_beat_prob, padded_downbeat_prob, mask): + """Function to compute the operations that must be computed piece by piece, and cannot be done in batch.""" + # unpad the predictions by truncating the padding positions + beat_prob = padded_beat_prob[mask] + downbeat_prob = padded_downbeat_prob[mask] + # build an artificial multiclass prediction, as suggested by Böck et al. + # again we limit the lower bound to avoid problems with the DBN + epsilon = 1e-5 + combined_act = np.vstack( + ( + np.maximum( + beat_prob.cpu().numpy() - downbeat_prob.cpu().numpy(), epsilon / 2 + ), + downbeat_prob.cpu().numpy(), + ) + ).T + # run the DBN + dbn_out = self.dbn(combined_act) + postp_beat = dbn_out[:, 0] + postp_downbeat = dbn_out[dbn_out[:, 1] == 1][:, 0] + return postp_beat, postp_downbeat + + +def deduplicate_peaks(peaks, width=1) -> np.ndarray: + """ + Replaces groups of adjacent peak frame indices that are each not more + than `width` frames apart by the average of the frame indices. + """ + result = [] + peaks = map(int, peaks) # ensure we get ordinary Python int objects + try: + p = next(peaks) + except StopIteration: + return np.array(result) + c = 1 + for p2 in peaks: + if p2 - p <= width: + c += 1 + p += (p2 - p) / c # update mean + else: + result.append(p) + p = p2 + c = 1 + result.append(p) + return np.array(result) diff --git a/beat_this/model/roformer.py b/beat_this/model/roformer.py new file mode 100644 index 0000000000000000000000000000000000000000..d1cdd3cef8509dedd0504207d68c2e1d284a5b58 --- /dev/null +++ b/beat_this/model/roformer.py @@ -0,0 +1,181 @@ +""" +Transformer with rotary position embedding, adapted from Phil Wang's repository +at https://github.com/lucidrains/BS-RoFormer (under MIT License). +""" + +import torch +import torch.nn.functional as F +from einops import rearrange +from torch import nn +from torch.nn import Module, ModuleList + +# helper functions + + +def exists(val): + return val is not None + + +# norm + + +class RMSNorm(Module): + def __init__(self, size, dim=-1): + super().__init__() + self.scale = size**0.5 + if dim >= 0: + raise ValueError(f"dim must be negative, got {dim}") + self.gamma = nn.Parameter(torch.ones((size,) + (1,) * (abs(dim) - 1))) + self.dim = dim + + def forward(self, x): + return F.normalize(x, dim=self.dim) * self.scale * self.gamma + + +# feedforward + + +class FeedForward(Module): + def __init__( + self, + dim, + mult=4, + dropout=0.0, + dim_out=None, + ): + super().__init__() + if dim_out is None: + dim_out = dim + dim_inner = int(dim * mult) + self.activation = nn.GELU() + self.net = nn.Sequential( + RMSNorm(dim), + nn.Linear(dim, dim_inner), + self.activation, + nn.Dropout(dropout), + nn.Linear(dim_inner, dim_out), + nn.Dropout(dropout), + ) + + def forward(self, x): + return self.net(x) + + +# attention + + +class Attend(nn.Module): + def __init__(self, dropout=0.0, scale=None): + super().__init__() + self.dropout = dropout + self.scale = scale + + def forward(self, q, k, v): + if exists(self.scale): + default_scale = q.shape[-1] ** -0.5 + q = q * (self.scale / default_scale) + + return F.scaled_dot_product_attention( + q, k, v, dropout_p=self.dropout if self.training else 0.0 + ) + + +class Attention(Module): + def __init__( + self, + dim, + heads=8, + dim_head=64, + dropout=0.0, + rotary_embed=None, + gating=True, + ): + super().__init__() + self.heads = heads + self.scale = dim_head**-0.5 + dim_inner = heads * dim_head + + self.rotary_embed = rotary_embed + + self.attend = Attend(dropout=dropout) + + self.norm = RMSNorm(dim) + self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False) + + if gating: + self.to_gates = nn.Linear(dim, heads) + else: + self.to_gates = None + + self.to_out = nn.Sequential( + nn.Linear(dim_inner, dim, bias=False), nn.Dropout(dropout) + ) + + def forward(self, x): + x = self.norm(x) + + q, k, v = rearrange( + self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads + ) + + if exists(self.rotary_embed): + q = self.rotary_embed.rotate_queries_or_keys(q) + k = self.rotary_embed.rotate_queries_or_keys(k) + + out = self.attend(q, k, v) + + if exists(self.to_gates): + gates = self.to_gates(x) + out = out * rearrange(gates, "b n h -> b h n 1").sigmoid() + + out = rearrange(out, "b h n d -> b n (h d)") + return self.to_out(out) + + +# Roformer + + +class Transformer(Module): + def __init__( + self, + *, + dim, + depth, + dim_head=32, + heads=16, + attn_dropout=0.1, + ff_dropout=0.1, + ff_mult=4, + norm_output=True, + rotary_embed=None, + gating=True, + ): + super().__init__() + self.layers = ModuleList([]) + + for _ in range(depth): + ff = FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout) + self.layers.append( + ModuleList( + [ + Attention( + dim=dim, + dim_head=dim_head, + heads=heads, + dropout=attn_dropout, + rotary_embed=rotary_embed, + gating=gating, + ), + ff, + ] + ) + ) + + self.norm = RMSNorm(dim) if norm_output else nn.Identity() + + def forward(self, x): + for attn, ff in self.layers: + x = attn(x) + x + x = ff(x) + x + x = self.norm(x) + return x diff --git a/beat_this/preprocessing.py b/beat_this/preprocessing.py new file mode 100644 index 0000000000000000000000000000000000000000..66ddcc4e24d3f55b38039616b259a52cc14f0d8f --- /dev/null +++ b/beat_this/preprocessing.py @@ -0,0 +1,59 @@ +import numpy as np +import torch +import torchaudio + + +def load_audio(path, dtype="float64"): + try: + waveform, samplerate = torchaudio.load(path, channels_first=False) + waveform = np.asanyarray(waveform.squeeze().numpy(), dtype=dtype) + return waveform, samplerate + except Exception: + # in case torchaudio fails, try soundfile + try: + import soundfile as sf + + return sf.read(path, dtype=dtype) + except Exception: + # some files are not readable by soundfile, try madmom + try: + import madmom + + return madmom.io.load_audio_file(str(path), dtype=dtype) + except Exception: + raise RuntimeError(f'Could not load audio from "{path}".') + + +class LogMelSpect(torch.nn.Module): + def __init__( + self, + sample_rate=22050, + n_fft=1024, + hop_length=441, + f_min=30, + f_max=11000, + n_mels=128, + mel_scale="slaney", + normalized="frame_length", + power=1, + log_multiplier=1000, + device="cpu", + ): + super().__init__() + self.spect_class = torchaudio.transforms.MelSpectrogram( + sample_rate=sample_rate, + n_fft=n_fft, + hop_length=hop_length, + f_min=f_min, + f_max=f_max, + n_mels=n_mels, + mel_scale=mel_scale, + normalized=normalized, + power=power, + ).to(device) + self.log_multiplier = log_multiplier + + def forward(self, x): + """Input is a waveform as a monodimensional array of shape T, + output is a 2D log mel spectrogram of shape (F,128).""" + return torch.log1p(self.log_multiplier * self.spect_class(x).T) diff --git a/beat_this/utils.py b/beat_this/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a364d4582481562335a495d1123595e42b0eafb9 --- /dev/null +++ b/beat_this/utils.py @@ -0,0 +1,111 @@ +from itertools import chain +from pathlib import Path + +import numpy as np + + +def index_to_framewise(index, length): + """Convert an index to a framewise sequence""" + sequence = np.zeros(length, dtype=bool) + sequence[index] = True + return sequence + + +def filename_to_augmentation(filename): + """Convert a filename to an augmentation factor.""" + parts = Path(filename).stem.split("_") + augmentations = {} + for part in parts[1:]: + if part.startswith("ps"): + augmentations["shift"] = int(part[2:]) + elif part.startswith("ts"): + augmentations["stretch"] = int(part[2:]) + return augmentations + + +def infer_beat_numbers(beats: np.ndarray, downbeats: np.ndarray) -> np.ndarray: + """ + From beat and downbeat times, infer a number for each beat such that each downbeat + is associated with a 1 and beats in between are counted upwards. + The function requires that all downbeats are also listed as beats. + + Args: + beats (numpy.ndarray): Array of beat positions in seconds (including downbeats). + downbeats (numpy.ndarray): Array of downbeat positions in seconds. + + Returns: + numbers (numpy.ndarray): Array of integer beat numbers. + """ + # check if all downbeats are beats + if not np.all(np.isin(downbeats, beats)): + raise ValueError("Not all downbeats are beats.") + + # handle pickup measure, by considering the beat count of the first full measure + if len(downbeats) >= 2: + # find the number of beats between the first two downbeats + first_downbeat, second_downbeat = np.searchsorted(beats, downbeats[:2]) + beats_in_first_measure = second_downbeat - first_downbeat + # find the number of beats before the first downbeat + pickup_beats = first_downbeat + # derive where to start counting + if pickup_beats < beats_in_first_measure: + start_counter = beats_in_first_measure - pickup_beats + else: + print( + "WARNING: There are more beats in the pickup measure than in the first measure. The beat count will start from 2 without trying to estimate the length of the pickup measure." + ) + start_counter = 1 + else: + print( + "WARNING: There are less than two downbeats in the predictions. Something may be wrong. The beat count will start from 2 without trying to estimate the length of the pickup measure." + ) + start_counter = 1 + + # assemble the beat numbers + numbers = [] + counter = start_counter + downbeats = chain(downbeats, [-1]) + next_downbeat = next(downbeats) + for beat in beats: + if beat == next_downbeat: + counter = 1 + next_downbeat = next(downbeats) + else: + counter += 1 + numbers.append(counter) + return np.asarray(numbers) + + +def save_beat_tsv(beats: np.ndarray, downbeats: np.ndarray, outpath: str) -> None: + """ + Save beat information to a tab-separated file in the standard .beats format: + each line has a time in seconds, a tab, and a beat number (1 = downbeat). + The function requires that all downbeats are also listed as beats. + + Args: + beats (numpy.ndarray): Array of beat positions in seconds (including downbeats). + downbeats (numpy.ndarray): Array of downbeat positions in seconds. + outpath (str): Path to the output TSV file. + + Returns: + None + """ + # infer beat numbers + numbers = infer_beat_numbers(beats, downbeats) + + # write the beat file + Path(outpath).parent.mkdir(parents=True, exist_ok=True) + try: + with open(outpath, "w") as f: + f.writelines(f"{beat}\t{number}\n" for beat, number in zip(beats, numbers)) + except KeyboardInterrupt: + outpath.unlink() # avoid half-written files + + +def replace_state_dict_key(state_dict: dict, old: str, new: str): + """Replaces `old` in all keys of `state_dict` with `new`.""" + keys = list(state_dict.keys()) # take snapshot of the keys + for key in keys: + if old in key: + state_dict[key.replace(old, new)] = state_dict.pop(key) + return state_dict diff --git a/model.json b/model.json new file mode 100644 index 0000000000000000000000000000000000000000..1801c6bdf2758b553829fe573815315d6dcdf807 --- /dev/null +++ b/model.json @@ -0,0 +1,18 @@ +{ + "name": "loopgen", + "package_dir": "src", + "entry_point": ["src.loopgen.LoopGen", "src.audiocraft.models.musicgen.MusicGen"], + "notes": "Multi-package model. (1) src/loopgen.py + src/lm_loopgen.py are LoopGen's own MAGNeT wrapper, copied from loopgen-repo/src/. (2) src/audiocraft/ is audiocraft v1.3.0 (loopgen-repo's own pin) vendored from source rather than pip-installed, because PyPI audiocraft==1.3.0 hard-pins torch==2.1.0/av==11.0.0/xformers<0.0.23, which conflicts with the torch 2.8.0 ZeroGPU needs -- vendoring sidesteps that pin entirely (same pattern already used by teamup-tech/instruct-MusicGen). All internal absolute 'from audiocraft.X' imports were mechanically rewritten to 'from src.audiocraft.X'; relative imports untouched. models/__init__.py and data/__init__.py were trimmed to drop AudioGen/MultiBandDiffusion/DiffusionUnet/DAC and training-dataset classes (unused, and multibanddiffusion pulls in the training-only solvers package). loaders.py's three torch.load calls got weights_only=False (MAGNeT's checkpoint bundles an OmegaConf config, not just tensors). (3) beat_this/ is vendored separately from ~/Documents/GitHub/beat-this (NOT loopgen's own git+ pin of CPJKU/beat_this) -- same upstream model, already proven on torch 2.11.0 via teamup-tech's own beat-this Space. See SOURCES.md.", + "additional_vendored_packages": [ + { + "package_dir": "beat_this", + "entry_point": "beat_this.inference.File2Beats", + "vendored_from": "~/Documents/GitHub/beat-this/beat_this/ (local clone, not loopgen-repo)" + } + ], + "checkpoints": { + "musicgen-medium": "auto-downloaded from facebook/musicgen-medium via audiocraft's MusicGen.get_pretrained() -- no manual hosting", + "magnet-medium-30secs": "auto-downloaded from facebook/magnet-medium-30secs via audiocraft's LoopGen.get_pretrained() -- no manual hosting", + "beat_this final0.ckpt": {"repo": "teamup-tech/beat-this", "repo_type": "space", "filename": "final0.ckpt", "size_mb": 81, "note": "reused from the already-deployed teamup-tech/beat-this Space via hf_hub_download(repo_type=\"space\"), not re-uploaded"} + } +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..9b94b1b71e19b8681e942d1da3987e60cefc9eaf --- /dev/null +++ b/requirements.txt @@ -0,0 +1,22 @@ +git+https://github.com/TEAMuP-dev/pyharp.git@develop +# model-specific deps below: +torch==2.8.0 +torchaudio==2.8.0 +av>=12,<13 +xformers==0.0.32.post2 +flashy +julius +numpy +tqdm +omegaconf +einops +num2words +transformers +sentencepiece +protobuf +spacy +huggingface_hub +soundfile +librosa +soxr +rotary-embedding-torch diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/audiocraft/__init__.py b/src/audiocraft/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..062ab7e514538f1dd48607973e90e4f708e72537 --- /dev/null +++ b/src/audiocraft/__init__.py @@ -0,0 +1,26 @@ +# 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. +""" +AudioCraft is a general framework for training audio generative models. +At the moment we provide the training code for: + +- [MusicGen](https://arxiv.org/abs/2306.05284), a state-of-the-art + text-to-music and melody+text autoregressive generative model. + For the solver, see `audiocraft.solvers.musicgen.MusicGenSolver`, and for the model, + `audiocraft.models.musicgen.MusicGen`. +- [AudioGen](https://arxiv.org/abs/2209.15352), a state-of-the-art + text-to-general-audio generative model. +- [EnCodec](https://arxiv.org/abs/2210.13438), efficient and high fidelity + neural audio codec which provides an excellent tokenizer for autoregressive language models. + See `audiocraft.solvers.compression.CompressionSolver`, and `audiocraft.models.encodec.EncodecModel`. +- [MultiBandDiffusion](TODO), alternative diffusion-based decoder compatible with EnCodec that + improves the perceived quality and reduces the artifacts coming from adversarial decoders. +""" + +# flake8: noqa +from . import data, modules, models + +__version__ = '1.3.0' diff --git a/src/audiocraft/data/__init__.py b/src/audiocraft/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..948589c8ec870f49078bfe1e43dd959780347b2b --- /dev/null +++ b/src/audiocraft/data/__init__.py @@ -0,0 +1,12 @@ +# 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. +"""Audio loading and writing support. Datasets for raw audio +or also including some metadata.""" + +# flake8: noqa +# dropped training-dataset classes (audio_dataset, info_audio_dataset, music_dataset, +# sound_dataset) to avoid unnecessary dependencies +from . import audio diff --git a/src/audiocraft/data/audio.py b/src/audiocraft/data/audio.py new file mode 100644 index 0000000000000000000000000000000000000000..a35dfd9c5a74330cfc3e1ba68932fffcb7c29750 --- /dev/null +++ b/src/audiocraft/data/audio.py @@ -0,0 +1,231 @@ +# 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. + +""" +Audio IO methods are defined in this module (info, read, write), +We rely on av library for faster read when possible, otherwise on torchaudio. +""" + +from dataclasses import dataclass +from pathlib import Path +import logging +import typing as tp + +import numpy as np +import soundfile +import torch +from torch.nn import functional as F + +import av +import subprocess as sp + +from .audio_utils import f32_pcm, normalize_audio + + +_av_initialized = False + + +def _init_av(): + global _av_initialized + if _av_initialized: + return + logger = logging.getLogger('libav.mp3') + logger.setLevel(logging.ERROR) + _av_initialized = True + + +@dataclass(frozen=True) +class AudioFileInfo: + sample_rate: int + duration: float + channels: int + + +def _av_info(filepath: tp.Union[str, Path]) -> AudioFileInfo: + _init_av() + with av.open(str(filepath)) as af: + stream = af.streams.audio[0] + sample_rate = stream.codec_context.sample_rate + duration = float(stream.duration * stream.time_base) + channels = stream.channels + return AudioFileInfo(sample_rate, duration, channels) + + +def _soundfile_info(filepath: tp.Union[str, Path]) -> AudioFileInfo: + info = soundfile.info(filepath) + return AudioFileInfo(info.samplerate, info.duration, info.channels) + + +def audio_info(filepath: tp.Union[str, Path]) -> AudioFileInfo: + # torchaudio no longer returns useful duration informations for some formats like mp3s. + filepath = Path(filepath) + if filepath.suffix in ['.flac', '.ogg']: # TODO: Validate .ogg can be safely read with av_info + # ffmpeg has some weird issue with flac. + return _soundfile_info(filepath) + else: + return _av_info(filepath) + + +def _av_read(filepath: tp.Union[str, Path], seek_time: float = 0, duration: float = -1.) -> tp.Tuple[torch.Tensor, int]: + """FFMPEG-based audio file reading using PyAV bindings. + Soundfile cannot read mp3 and av_read is more efficient than torchaudio. + + Args: + filepath (str or Path): Path to audio file to read. + seek_time (float): Time at which to start reading in the file. + duration (float): Duration to read from the file. If set to -1, the whole file is read. + Returns: + tuple of torch.Tensor, int: Tuple containing audio data and sample rate + """ + _init_av() + with av.open(str(filepath)) as af: + stream = af.streams.audio[0] + sr = stream.codec_context.sample_rate + num_frames = int(sr * duration) if duration >= 0 else -1 + frame_offset = int(sr * seek_time) + # we need a small negative offset otherwise we get some edge artifact + # from the mp3 decoder. + af.seek(int(max(0, (seek_time - 0.1)) / stream.time_base), stream=stream) + frames = [] + length = 0 + for frame in af.decode(streams=stream.index): + current_offset = int(frame.rate * frame.pts * frame.time_base) + strip = max(0, frame_offset - current_offset) + buf = torch.from_numpy(frame.to_ndarray()) + if buf.shape[0] != stream.channels: + buf = buf.view(-1, stream.channels).t() + buf = buf[:, strip:] + frames.append(buf) + length += buf.shape[1] + if num_frames > 0 and length >= num_frames: + break + assert frames + # If the above assert fails, it is likely because we seeked past the end of file point, + # in which case ffmpeg returns a single frame with only zeros, and a weird timestamp. + # This will need proper debugging, in due time. + wav = torch.cat(frames, dim=1) + assert wav.shape[0] == stream.channels + if num_frames > 0: + wav = wav[:, :num_frames] + return f32_pcm(wav), sr + + +def audio_read(filepath: tp.Union[str, Path], seek_time: float = 0., + duration: float = -1., pad: bool = False) -> tp.Tuple[torch.Tensor, int]: + """Read audio by picking the most appropriate backend tool based on the audio format. + + Args: + filepath (str or Path): Path to audio file to read. + seek_time (float): Time at which to start reading in the file. + duration (float): Duration to read from the file. If set to -1, the whole file is read. + pad (bool): Pad output audio if not reaching expected duration. + Returns: + tuple of torch.Tensor, int: Tuple containing audio data and sample rate. + """ + fp = Path(filepath) + if fp.suffix in ['.flac', '.ogg']: # TODO: check if we can safely use av_read for .ogg + # There is some bug with ffmpeg and reading flac + info = _soundfile_info(filepath) + frames = -1 if duration <= 0 else int(duration * info.sample_rate) + frame_offset = int(seek_time * info.sample_rate) + wav, sr = soundfile.read(filepath, start=frame_offset, frames=frames, dtype=np.float32) + assert info.sample_rate == sr, f"Mismatch of sample rates {info.sample_rate} {sr}" + wav = torch.from_numpy(wav).t().contiguous() + if len(wav.shape) == 1: + wav = torch.unsqueeze(wav, 0) + else: + wav, sr = _av_read(filepath, seek_time, duration) + if pad and duration > 0: + expected_frames = int(duration * sr) + wav = F.pad(wav, (0, expected_frames - wav.shape[-1])) + return wav, sr + + +def _piping_to_ffmpeg(out_path: tp.Union[str, Path], wav: torch.Tensor, sample_rate: int, flags: tp.List[str]): + # ffmpeg is always installed and torchaudio is a bit unstable lately, so let's bypass it entirely. + assert wav.dim() == 2, wav.shape + command = [ + 'ffmpeg', + '-loglevel', 'error', + '-y', '-f', 'f32le', '-ar', str(sample_rate), '-ac', str(wav.shape[0]), + '-i', '-'] + flags + [str(out_path)] + input_ = f32_pcm(wav).t().detach().cpu().numpy().tobytes() + sp.run(command, input=input_, check=True) + + +def audio_write(stem_name: tp.Union[str, Path], + wav: torch.Tensor, sample_rate: int, + format: str = 'wav', mp3_rate: int = 320, ogg_rate: tp.Optional[int] = None, + normalize: bool = True, strategy: str = 'peak', peak_clip_headroom_db: float = 1, + rms_headroom_db: float = 18, loudness_headroom_db: float = 14, + loudness_compressor: bool = False, + log_clipping: bool = True, make_parent_dir: bool = True, + add_suffix: bool = True) -> Path: + """Convenience function for saving audio to disk. Returns the filename the audio was written to. + + Args: + stem_name (str or Path): Filename without extension which will be added automatically. + wav (torch.Tensor): Audio data to save. + sample_rate (int): Sample rate of audio data. + format (str): Either "wav", "mp3", "ogg", or "flac". + mp3_rate (int): kbps when using mp3s. + ogg_rate (int): kbps when using ogg/vorbis. If not provided, let ffmpeg decide for itself. + normalize (bool): if `True` (default), normalizes according to the prescribed + strategy (see after). If `False`, the strategy is only used in case clipping + would happen. + strategy (str): Can be either 'clip', 'peak', or 'rms'. Default is 'peak', + i.e. audio is normalized by its largest value. RMS normalizes by root-mean-square + with extra headroom to avoid clipping. 'clip' just clips. + peak_clip_headroom_db (float): Headroom in dB when doing 'peak' or 'clip' strategy. + rms_headroom_db (float): Headroom in dB when doing 'rms' strategy. This must be much larger + than the `peak_clip` one to avoid further clipping. + loudness_headroom_db (float): Target loudness for loudness normalization. + loudness_compressor (bool): Uses tanh for soft clipping when strategy is 'loudness'. + when strategy is 'loudness' log_clipping (bool): If True, basic logging on stderr when clipping still + occurs despite strategy (only for 'rms'). + make_parent_dir (bool): Make parent directory if it doesn't exist. + Returns: + Path: Path of the saved audio. + """ + assert wav.dtype.is_floating_point, "wav is not floating point" + if wav.dim() == 1: + wav = wav[None] + elif wav.dim() > 2: + raise ValueError("Input wav should be at most 2 dimension.") + assert wav.isfinite().all() + wav = normalize_audio(wav, normalize, strategy, peak_clip_headroom_db, + rms_headroom_db, loudness_headroom_db, loudness_compressor, + log_clipping=log_clipping, sample_rate=sample_rate, + stem_name=str(stem_name)) + if format == 'mp3': + suffix = '.mp3' + flags = ['-f', 'mp3', '-c:a', 'libmp3lame', '-b:a', f'{mp3_rate}k'] + elif format == 'wav': + suffix = '.wav' + flags = ['-f', 'wav', '-c:a', 'pcm_s16le'] + elif format == 'ogg': + suffix = '.ogg' + flags = ['-f', 'ogg', '-c:a', 'libvorbis'] + if ogg_rate is not None: + flags += ['-b:a', f'{ogg_rate}k'] + elif format == 'flac': + suffix = '.flac' + flags = ['-f', 'flac'] + else: + raise RuntimeError(f"Invalid format {format}. Only wav or mp3 are supported.") + if not add_suffix: + suffix = '' + path = Path(str(stem_name) + suffix) + if make_parent_dir: + path.parent.mkdir(exist_ok=True, parents=True) + try: + _piping_to_ffmpeg(path, wav, sample_rate, flags) + except Exception: + if path.exists(): + # we do not want to leave half written files around. + path.unlink() + raise + return path diff --git a/src/audiocraft/data/audio_dataset.py b/src/audiocraft/data/audio_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..9d7442526186b3712f5d4754f928a40ecd964174 --- /dev/null +++ b/src/audiocraft/data/audio_dataset.py @@ -0,0 +1,587 @@ +# 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. +"""AudioDataset support. In order to handle a larger number of files +without having to scan again the folders, we precompute some metadata +(filename, sample rate, duration), and use that to efficiently sample audio segments. +""" +import argparse +import copy +from concurrent.futures import ThreadPoolExecutor, Future +from dataclasses import dataclass, fields +from contextlib import ExitStack +from functools import lru_cache +import gzip +import json +import logging +import os +from pathlib import Path +import random +import sys +import typing as tp + +import torch +import torch.nn.functional as F + +from .audio import audio_read, audio_info +from .audio_utils import convert_audio +from .zip import PathInZip + +try: + import dora +except ImportError: + dora = None # type: ignore + + +@dataclass(order=True) +class BaseInfo: + + @classmethod + def _dict2fields(cls, dictionary: dict): + return { + field.name: dictionary[field.name] + for field in fields(cls) if field.name in dictionary + } + + @classmethod + def from_dict(cls, dictionary: dict): + _dictionary = cls._dict2fields(dictionary) + return cls(**_dictionary) + + def to_dict(self): + return { + field.name: self.__getattribute__(field.name) + for field in fields(self) + } + + +@dataclass(order=True) +class AudioMeta(BaseInfo): + path: str + duration: float + sample_rate: int + amplitude: tp.Optional[float] = None + weight: tp.Optional[float] = None + # info_path is used to load additional information about the audio file that is stored in zip files. + info_path: tp.Optional[PathInZip] = None + + @classmethod + def from_dict(cls, dictionary: dict): + base = cls._dict2fields(dictionary) + if 'info_path' in base and base['info_path'] is not None: + base['info_path'] = PathInZip(base['info_path']) + return cls(**base) + + def to_dict(self): + d = super().to_dict() + if d['info_path'] is not None: + d['info_path'] = str(d['info_path']) + return d + + +@dataclass(order=True) +class SegmentInfo(BaseInfo): + meta: AudioMeta + seek_time: float + # The following values are given once the audio is processed, e.g. + # at the target sample rate and target number of channels. + n_frames: int # actual number of frames without padding + total_frames: int # total number of frames, padding included + sample_rate: int # actual sample rate + channels: int # number of audio channels. + + +DEFAULT_EXTS = ['.wav', '.mp3', '.flac', '.ogg', '.m4a'] + +logger = logging.getLogger(__name__) + + +def _get_audio_meta(file_path: str, minimal: bool = True) -> AudioMeta: + """AudioMeta from a path to an audio file. + + Args: + file_path (str): Resolved path of valid audio file. + minimal (bool): Whether to only load the minimal set of metadata (takes longer if not). + Returns: + AudioMeta: Audio file path and its metadata. + """ + info = audio_info(file_path) + amplitude: tp.Optional[float] = None + if not minimal: + wav, sr = audio_read(file_path) + amplitude = wav.abs().max().item() + return AudioMeta(file_path, info.duration, info.sample_rate, amplitude) + + +def _resolve_audio_meta(m: AudioMeta, fast: bool = True) -> AudioMeta: + """If Dora is available as a dependency, try to resolve potential relative paths + in list of AudioMeta. This method is expected to be used when loading meta from file. + + Args: + m (AudioMeta): Audio meta to resolve. + fast (bool): If True, uses a really fast check for determining if a file + is already absolute or not. Only valid on Linux/Mac. + Returns: + AudioMeta: Audio meta with resolved path. + """ + def is_abs(m): + if fast: + return str(m)[0] == '/' + else: + os.path.isabs(str(m)) + + if not dora: + return m + + if not is_abs(m.path): + m.path = dora.git_save.to_absolute_path(m.path) + if m.info_path is not None and not is_abs(m.info_path.zip_path): + m.info_path.zip_path = dora.git_save.to_absolute_path(m.path) + return m + + +def find_audio_files(path: tp.Union[Path, str], + exts: tp.List[str] = DEFAULT_EXTS, + resolve: bool = True, + minimal: bool = True, + progress: bool = False, + workers: int = 0) -> tp.List[AudioMeta]: + """Build a list of AudioMeta from a given path, + collecting relevant audio files and fetching meta info. + + Args: + path (str or Path): Path to folder containing audio files. + exts (list of str): List of file extensions to consider for audio files. + minimal (bool): Whether to only load the minimal set of metadata (takes longer if not). + progress (bool): Whether to log progress on audio files collection. + workers (int): number of parallel workers, if 0, use only the current thread. + Returns: + list of AudioMeta: List of audio file path and its metadata. + """ + audio_files = [] + futures: tp.List[Future] = [] + pool: tp.Optional[ThreadPoolExecutor] = None + with ExitStack() as stack: + if workers > 0: + pool = ThreadPoolExecutor(workers) + stack.enter_context(pool) + + if progress: + print("Finding audio files...") + for root, folders, files in os.walk(path, followlinks=True): + for file in files: + full_path = Path(root) / file + if full_path.suffix.lower() in exts: + audio_files.append(full_path) + if pool is not None: + futures.append(pool.submit(_get_audio_meta, str(audio_files[-1]), minimal)) + if progress: + print(format(len(audio_files), " 8d"), end='\r', file=sys.stderr) + + if progress: + print("Getting audio metadata...") + meta: tp.List[AudioMeta] = [] + for idx, file_path in enumerate(audio_files): + try: + if pool is None: + m = _get_audio_meta(str(file_path), minimal) + else: + m = futures[idx].result() + if resolve: + m = _resolve_audio_meta(m) + except Exception as err: + print("Error with", str(file_path), err, file=sys.stderr) + continue + meta.append(m) + if progress: + print(format((1 + idx) / len(audio_files), " 3.1%"), end='\r', file=sys.stderr) + meta.sort() + return meta + + +def load_audio_meta(path: tp.Union[str, Path], + resolve: bool = True, fast: bool = True) -> tp.List[AudioMeta]: + """Load list of AudioMeta from an optionally compressed json file. + + Args: + path (str or Path): Path to JSON file. + resolve (bool): Whether to resolve the path from AudioMeta (default=True). + fast (bool): activates some tricks to make things faster. + Returns: + list of AudioMeta: List of audio file path and its total duration. + """ + open_fn = gzip.open if str(path).lower().endswith('.gz') else open + with open_fn(path, 'rb') as fp: # type: ignore + lines = fp.readlines() + meta = [] + for line in lines: + d = json.loads(line) + m = AudioMeta.from_dict(d) + if resolve: + m = _resolve_audio_meta(m, fast=fast) + meta.append(m) + return meta + + +def save_audio_meta(path: tp.Union[str, Path], meta: tp.List[AudioMeta]): + """Save the audio metadata to the file pointer as json. + + Args: + path (str or Path): Path to JSON file. + metadata (list of BaseAudioMeta): List of audio meta to save. + """ + Path(path).parent.mkdir(exist_ok=True, parents=True) + open_fn = gzip.open if str(path).lower().endswith('.gz') else open + with open_fn(path, 'wb') as fp: # type: ignore + for m in meta: + json_str = json.dumps(m.to_dict()) + '\n' + json_bytes = json_str.encode('utf-8') + fp.write(json_bytes) + + +class AudioDataset: + """Base audio dataset. + + The dataset takes a list of AudioMeta and create a dataset composed of segments of audio + and potentially additional information, by creating random segments from the list of audio + files referenced in the metadata and applying minimal data pre-processing such as resampling, + mixing of channels, padding, etc. + + If no segment_duration value is provided, the AudioDataset will return the full wav for each + audio file. Otherwise, it will randomly sample audio files and create a segment of the specified + duration, applying padding if required. + + By default, only the torch Tensor corresponding to the waveform is returned. Setting return_info=True + allows to return a tuple containing the torch Tensor and additional metadata on the segment and the + original audio meta. + + Note that you can call `start_epoch(epoch)` in order to get + a deterministic "randomization" for `shuffle=True`. + For a given epoch and dataset index, this will always return the same extract. + You can get back some diversity by setting the `shuffle_seed` param. + + Args: + meta (list of AudioMeta): List of audio files metadata. + segment_duration (float, optional): Optional segment duration of audio to load. + If not specified, the dataset will load the full audio segment from the file. + shuffle (bool): Set to `True` to have the data reshuffled at every epoch. + sample_rate (int): Target sample rate of the loaded audio samples. + channels (int): Target number of channels of the loaded audio samples. + sample_on_duration (bool): Set to `True` to sample segments with probability + dependent on audio file duration. This is only used if `segment_duration` is provided. + sample_on_weight (bool): Set to `True` to sample segments using the `weight` entry of + `AudioMeta`. If `sample_on_duration` is also True, the actual weight will be the product + of the file duration and file weight. This is only used if `segment_duration` is provided. + min_segment_ratio (float): Minimum segment ratio to use when the audio file + is shorter than the desired segment. + max_read_retry (int): Maximum number of retries to sample an audio segment from the dataset. + return_info (bool): Whether to return the wav only or return wav along with segment info and metadata. + min_audio_duration (float, optional): Minimum audio file duration, in seconds, if provided + audio shorter than this will be filtered out. + max_audio_duration (float, optional): Maximal audio file duration in seconds, if provided + audio longer than this will be filtered out. + shuffle_seed (int): can be used to further randomize + load_wav (bool): if False, skip loading the wav but returns a tensor of 0 + with the expected segment_duration (which must be provided if load_wav is False). + permutation_on_files (bool): only if `sample_on_weight` and `sample_on_duration` + are False. Will ensure a permutation on files when going through the dataset. + In that case the epoch number must be provided in order for the model + to continue the permutation across epochs. In that case, it is assumed + that `num_samples = total_batch_size * num_updates_per_epoch`, with + `total_batch_size` the overall batch size accounting for all gpus. + """ + def __init__(self, + meta: tp.List[AudioMeta], + segment_duration: tp.Optional[float] = None, + shuffle: bool = True, + num_samples: int = 10_000, + sample_rate: int = 48_000, + channels: int = 2, + pad: bool = True, + sample_on_duration: bool = True, + sample_on_weight: bool = True, + min_segment_ratio: float = 0.5, + max_read_retry: int = 10, + return_info: bool = False, + min_audio_duration: tp.Optional[float] = None, + max_audio_duration: tp.Optional[float] = None, + shuffle_seed: int = 0, + load_wav: bool = True, + permutation_on_files: bool = False, + ): + assert len(meta) > 0, "No audio meta provided to AudioDataset. Please check loading of audio meta." + assert segment_duration is None or segment_duration > 0 + assert segment_duration is None or min_segment_ratio >= 0 + self.segment_duration = segment_duration + self.min_segment_ratio = min_segment_ratio + self.max_audio_duration = max_audio_duration + self.min_audio_duration = min_audio_duration + if self.min_audio_duration is not None and self.max_audio_duration is not None: + assert self.min_audio_duration <= self.max_audio_duration + self.meta: tp.List[AudioMeta] = self._filter_duration(meta) + assert len(self.meta) # Fail fast if all data has been filtered. + self.total_duration = sum(d.duration for d in self.meta) + + if segment_duration is None: + num_samples = len(self.meta) + self.num_samples = num_samples + self.shuffle = shuffle + self.sample_rate = sample_rate + self.channels = channels + self.pad = pad + self.sample_on_weight = sample_on_weight + self.sample_on_duration = sample_on_duration + self.sampling_probabilities = self._get_sampling_probabilities() + self.max_read_retry = max_read_retry + self.return_info = return_info + self.shuffle_seed = shuffle_seed + self.current_epoch: tp.Optional[int] = None + self.load_wav = load_wav + if not load_wav: + assert segment_duration is not None + self.permutation_on_files = permutation_on_files + if permutation_on_files: + assert not self.sample_on_duration + assert not self.sample_on_weight + assert self.shuffle + + def start_epoch(self, epoch: int): + self.current_epoch = epoch + + def __len__(self): + return self.num_samples + + def _get_sampling_probabilities(self, normalized: bool = True): + """Return the sampling probabilities for each file inside `self.meta`.""" + scores: tp.List[float] = [] + for file_meta in self.meta: + score = 1. + if self.sample_on_weight and file_meta.weight is not None: + score *= file_meta.weight + if self.sample_on_duration: + score *= file_meta.duration + scores.append(score) + probabilities = torch.tensor(scores) + if normalized: + probabilities /= probabilities.sum() + return probabilities + + @staticmethod + @lru_cache(16) + def _get_file_permutation(num_files: int, permutation_index: int, base_seed: int): + # Used to keep the most recent files permutation in memory implicitely. + # will work unless someone is using a lot of Datasets in parallel. + rng = torch.Generator() + rng.manual_seed(base_seed + permutation_index) + return torch.randperm(num_files, generator=rng) + + def sample_file(self, index: int, rng: torch.Generator) -> AudioMeta: + """Sample a given file from `self.meta`. Can be overridden in subclasses. + This is only called if `segment_duration` is not None. + + You must use the provided random number generator `rng` for reproducibility. + You can further make use of the index accessed. + """ + if self.permutation_on_files: + assert self.current_epoch is not None + total_index = self.current_epoch * len(self) + index + permutation_index = total_index // len(self.meta) + relative_index = total_index % len(self.meta) + permutation = AudioDataset._get_file_permutation( + len(self.meta), permutation_index, self.shuffle_seed) + file_index = permutation[relative_index] + return self.meta[file_index] + + if not self.sample_on_weight and not self.sample_on_duration: + file_index = int(torch.randint(len(self.sampling_probabilities), (1,), generator=rng).item()) + else: + file_index = int(torch.multinomial(self.sampling_probabilities, 1, generator=rng).item()) + + return self.meta[file_index] + + def _audio_read(self, path: str, seek_time: float = 0, duration: float = -1): + # Override this method in subclass if needed. + if self.load_wav: + return audio_read(path, seek_time, duration, pad=False) + else: + assert self.segment_duration is not None + n_frames = int(self.sample_rate * self.segment_duration) + return torch.zeros(self.channels, n_frames), self.sample_rate + + def __getitem__(self, index: int) -> tp.Union[torch.Tensor, tp.Tuple[torch.Tensor, SegmentInfo]]: + if self.segment_duration is None: + file_meta = self.meta[index] + out, sr = audio_read(file_meta.path) + out = convert_audio(out, sr, self.sample_rate, self.channels) + n_frames = out.shape[-1] + segment_info = SegmentInfo(file_meta, seek_time=0., n_frames=n_frames, total_frames=n_frames, + sample_rate=self.sample_rate, channels=out.shape[0]) + else: + rng = torch.Generator() + if self.shuffle: + # We use index, plus extra randomness, either totally random if we don't know the epoch. + # otherwise we make use of the epoch number and optional shuffle_seed. + if self.current_epoch is None: + rng.manual_seed(index + self.num_samples * random.randint(0, 2**24)) + else: + rng.manual_seed(index + self.num_samples * (self.current_epoch + self.shuffle_seed)) + else: + # We only use index + rng.manual_seed(index) + + for retry in range(self.max_read_retry): + file_meta = self.sample_file(index, rng) + # We add some variance in the file position even if audio file is smaller than segment + # without ending up with empty segments + max_seek = max(0, file_meta.duration - self.segment_duration * self.min_segment_ratio) + seek_time = torch.rand(1, generator=rng).item() * max_seek + try: + out, sr = audio_read(file_meta.path, seek_time, self.segment_duration, pad=False) + out = convert_audio(out, sr, self.sample_rate, self.channels) + n_frames = out.shape[-1] + target_frames = int(self.segment_duration * self.sample_rate) + if self.pad: + out = F.pad(out, (0, target_frames - n_frames)) + segment_info = SegmentInfo(file_meta, seek_time, n_frames=n_frames, total_frames=target_frames, + sample_rate=self.sample_rate, channels=out.shape[0]) + except Exception as exc: + logger.warning("Error opening file %s: %r", file_meta.path, exc) + if retry == self.max_read_retry - 1: + raise + else: + break + + if self.return_info: + # Returns the wav and additional information on the wave segment + return out, segment_info + else: + return out + + def collater(self, samples): + """The collater function has to be provided to the dataloader + if AudioDataset has return_info=True in order to properly collate + the samples of a batch. + """ + if self.segment_duration is None and len(samples) > 1: + assert self.pad, "Must allow padding when batching examples of different durations." + + # In this case the audio reaching the collater is of variable length as segment_duration=None. + to_pad = self.segment_duration is None and self.pad + if to_pad: + max_len = max([wav.shape[-1] for wav, _ in samples]) + + def _pad_wav(wav): + return F.pad(wav, (0, max_len - wav.shape[-1])) + + if self.return_info: + if len(samples) > 0: + assert len(samples[0]) == 2 + assert isinstance(samples[0][0], torch.Tensor) + assert isinstance(samples[0][1], SegmentInfo) + + wavs = [wav for wav, _ in samples] + segment_infos = [copy.deepcopy(info) for _, info in samples] + + if to_pad: + # Each wav could be of a different duration as they are not segmented. + for i in range(len(samples)): + # Determines the total length of the signal with padding, so we update here as we pad. + segment_infos[i].total_frames = max_len + wavs[i] = _pad_wav(wavs[i]) + + wav = torch.stack(wavs) + return wav, segment_infos + else: + assert isinstance(samples[0], torch.Tensor) + if to_pad: + samples = [_pad_wav(s) for s in samples] + return torch.stack(samples) + + def _filter_duration(self, meta: tp.List[AudioMeta]) -> tp.List[AudioMeta]: + """Filters out audio files with audio durations that will not allow to sample examples from them.""" + orig_len = len(meta) + + # Filter data that is too short. + if self.min_audio_duration is not None: + meta = [m for m in meta if m.duration >= self.min_audio_duration] + + # Filter data that is too long. + if self.max_audio_duration is not None: + meta = [m for m in meta if m.duration <= self.max_audio_duration] + + filtered_len = len(meta) + removed_percentage = 100*(1-float(filtered_len)/orig_len) + msg = 'Removed %.2f percent of the data because it was too short or too long.' % removed_percentage + if removed_percentage < 10: + logging.debug(msg) + else: + logging.warning(msg) + return meta + + @classmethod + def from_meta(cls, root: tp.Union[str, Path], **kwargs): + """Instantiate AudioDataset from a path to a directory containing a manifest as a jsonl file. + + Args: + root (str or Path): Path to root folder containing audio files. + kwargs: Additional keyword arguments for the AudioDataset. + """ + root = Path(root) + if root.is_dir(): + if (root / 'data.jsonl').exists(): + root = root / 'data.jsonl' + elif (root / 'data.jsonl.gz').exists(): + root = root / 'data.jsonl.gz' + else: + raise ValueError("Don't know where to read metadata from in the dir. " + "Expecting either a data.jsonl or data.jsonl.gz file but none found.") + meta = load_audio_meta(root) + return cls(meta, **kwargs) + + @classmethod + def from_path(cls, root: tp.Union[str, Path], minimal_meta: bool = True, + exts: tp.List[str] = DEFAULT_EXTS, **kwargs): + """Instantiate AudioDataset from a path containing (possibly nested) audio files. + + Args: + root (str or Path): Path to root folder containing audio files. + minimal_meta (bool): Whether to only load minimal metadata or not. + exts (list of str): Extensions for audio files. + kwargs: Additional keyword arguments for the AudioDataset. + """ + root = Path(root) + if root.is_file(): + meta = load_audio_meta(root, resolve=True) + else: + meta = find_audio_files(root, exts, minimal=minimal_meta, resolve=True) + return cls(meta, **kwargs) + + +def main(): + logging.basicConfig(stream=sys.stderr, level=logging.INFO) + parser = argparse.ArgumentParser( + prog='audio_dataset', + description='Generate .jsonl files by scanning a folder.') + parser.add_argument('root', help='Root folder with all the audio files') + parser.add_argument('output_meta_file', + help='Output file to store the metadata, ') + parser.add_argument('--complete', + action='store_false', dest='minimal', default=True, + help='Retrieve all metadata, even the one that are expansive ' + 'to compute (e.g. normalization).') + parser.add_argument('--resolve', + action='store_true', default=False, + help='Resolve the paths to be absolute and with no symlinks.') + parser.add_argument('--workers', + default=10, type=int, + help='Number of workers.') + args = parser.parse_args() + meta = find_audio_files(args.root, DEFAULT_EXTS, progress=True, + resolve=args.resolve, minimal=args.minimal, workers=args.workers) + save_audio_meta(args.output_meta_file, meta) + + +if __name__ == '__main__': + main() diff --git a/src/audiocraft/data/audio_utils.py b/src/audiocraft/data/audio_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9d3129b84b114c5572078295604279884c79f2cc --- /dev/null +++ b/src/audiocraft/data/audio_utils.py @@ -0,0 +1,176 @@ +# 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. +"""Various utilities for audio convertion (pcm format, sample rate and channels), +and volume normalization.""" +import sys +import typing as tp + +import julius +import torch +import torchaudio + + +def convert_audio_channels(wav: torch.Tensor, channels: int = 2) -> torch.Tensor: + """Convert audio to the given number of channels. + + Args: + wav (torch.Tensor): Audio wave of shape [B, C, T]. + channels (int): Expected number of channels as output. + Returns: + torch.Tensor: Downmixed or unchanged audio wave [B, C, T]. + """ + *shape, src_channels, length = wav.shape + if src_channels == channels: + pass + elif channels == 1: + # Case 1: + # The caller asked 1-channel audio, and the stream has multiple + # channels, downmix all channels. + wav = wav.mean(dim=-2, keepdim=True) + elif src_channels == 1: + # Case 2: + # The caller asked for multiple channels, but the input file has + # a single channel, replicate the audio over all channels. + wav = wav.expand(*shape, channels, length) + elif src_channels >= channels: + # Case 3: + # The caller asked for multiple channels, and the input file has + # more channels than requested. In that case return the first channels. + wav = wav[..., :channels, :] + else: + # Case 4: What is a reasonable choice here? + raise ValueError('The audio file has less channels than requested but is not mono.') + return wav + + +def convert_audio(wav: torch.Tensor, from_rate: float, + to_rate: float, to_channels: int) -> torch.Tensor: + """Convert audio to new sample rate and number of audio channels.""" + wav = julius.resample_frac(wav, int(from_rate), int(to_rate)) + wav = convert_audio_channels(wav, to_channels) + return wav + + +def normalize_loudness(wav: torch.Tensor, sample_rate: int, loudness_headroom_db: float = 14, + loudness_compressor: bool = False, energy_floor: float = 2e-3): + """Normalize an input signal to a user loudness in dB LKFS. + Audio loudness is defined according to the ITU-R BS.1770-4 recommendation. + + Args: + wav (torch.Tensor): Input multichannel audio data. + sample_rate (int): Sample rate. + loudness_headroom_db (float): Target loudness of the output in dB LUFS. + loudness_compressor (bool): Uses tanh for soft clipping. + energy_floor (float): anything below that RMS level will not be rescaled. + Returns: + torch.Tensor: Loudness normalized output data. + """ + energy = wav.pow(2).mean().sqrt().item() + if energy < energy_floor: + return wav + transform = torchaudio.transforms.Loudness(sample_rate) + input_loudness_db = transform(wav).item() + # calculate the gain needed to scale to the desired loudness level + delta_loudness = -loudness_headroom_db - input_loudness_db + gain = 10.0 ** (delta_loudness / 20.0) + output = gain * wav + if loudness_compressor: + output = torch.tanh(output) + assert output.isfinite().all(), (input_loudness_db, wav.pow(2).mean().sqrt()) + return output + + +def _clip_wav(wav: torch.Tensor, log_clipping: bool = False, stem_name: tp.Optional[str] = None) -> None: + """Utility function to clip the audio with logging if specified.""" + max_scale = wav.abs().max() + if log_clipping and max_scale > 1: + clamp_prob = (wav.abs() > 1).float().mean().item() + print(f"CLIPPING {stem_name or ''} happening with proba (a bit of clipping is okay):", + clamp_prob, "maximum scale: ", max_scale.item(), file=sys.stderr) + wav.clamp_(-1, 1) + + +def normalize_audio(wav: torch.Tensor, normalize: bool = True, + strategy: str = 'peak', peak_clip_headroom_db: float = 1, + rms_headroom_db: float = 18, loudness_headroom_db: float = 14, + loudness_compressor: bool = False, log_clipping: bool = False, + sample_rate: tp.Optional[int] = None, + stem_name: tp.Optional[str] = None) -> torch.Tensor: + """Normalize the audio according to the prescribed strategy (see after). + + Args: + wav (torch.Tensor): Audio data. + normalize (bool): if `True` (default), normalizes according to the prescribed + strategy (see after). If `False`, the strategy is only used in case clipping + would happen. + strategy (str): Can be either 'clip', 'peak', or 'rms'. Default is 'peak', + i.e. audio is normalized by its largest value. RMS normalizes by root-mean-square + with extra headroom to avoid clipping. 'clip' just clips. + peak_clip_headroom_db (float): Headroom in dB when doing 'peak' or 'clip' strategy. + rms_headroom_db (float): Headroom in dB when doing 'rms' strategy. This must be much larger + than the `peak_clip` one to avoid further clipping. + loudness_headroom_db (float): Target loudness for loudness normalization. + loudness_compressor (bool): If True, uses tanh based soft clipping. + log_clipping (bool): If True, basic logging on stderr when clipping still + occurs despite strategy (only for 'rms'). + sample_rate (int): Sample rate for the audio data (required for loudness). + stem_name (str, optional): Stem name for clipping logging. + Returns: + torch.Tensor: Normalized audio. + """ + scale_peak = 10 ** (-peak_clip_headroom_db / 20) + scale_rms = 10 ** (-rms_headroom_db / 20) + if strategy == 'peak': + rescaling = (scale_peak / wav.abs().max()) + if normalize or rescaling < 1: + wav = wav * rescaling + elif strategy == 'clip': + wav = wav.clamp(-scale_peak, scale_peak) + elif strategy == 'rms': + mono = wav.mean(dim=0) + rescaling = scale_rms / mono.pow(2).mean().sqrt() + if normalize or rescaling < 1: + wav = wav * rescaling + _clip_wav(wav, log_clipping=log_clipping, stem_name=stem_name) + elif strategy == 'loudness': + assert sample_rate is not None, "Loudness normalization requires sample rate." + wav = normalize_loudness(wav, sample_rate, loudness_headroom_db, loudness_compressor) + _clip_wav(wav, log_clipping=log_clipping, stem_name=stem_name) + else: + assert wav.abs().max() < 1 + assert strategy == '' or strategy == 'none', f"Unexpected strategy: '{strategy}'" + return wav + + +def f32_pcm(wav: torch.Tensor) -> torch.Tensor: + """Convert audio to float 32 bits PCM format. + """ + if wav.dtype.is_floating_point: + return wav + elif wav.dtype == torch.int16: + return wav.float() / 2**15 + elif wav.dtype == torch.int32: + return wav.float() / 2**31 + raise ValueError(f"Unsupported wav dtype: {wav.dtype}") + + +def i16_pcm(wav: torch.Tensor) -> torch.Tensor: + """Convert audio to int 16 bits PCM format. + + ..Warning:: There exist many formula for doing this conversion. None are perfect + due to the asymmetry of the int16 range. One either have possible clipping, DC offset, + or inconsistencies with f32_pcm. If the given wav doesn't have enough headroom, + it is possible that `i16_pcm(f32_pcm)) != Identity`. + """ + if wav.dtype.is_floating_point: + assert wav.abs().max() <= 1 + candidate = (wav * 2 ** 15).round() + if candidate.max() >= 2 ** 15: # clipping would occur + candidate = (wav * (2 ** 15 - 1)).round() + return candidate.short() + else: + assert wav.dtype == torch.int16 + return wav diff --git a/src/audiocraft/data/zip.py b/src/audiocraft/data/zip.py new file mode 100644 index 0000000000000000000000000000000000000000..f0b17849d36991e7def35a14d3d518b9d867ce36 --- /dev/null +++ b/src/audiocraft/data/zip.py @@ -0,0 +1,76 @@ +# 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. +"""Utility for reading some info from inside a zip file. +""" + +import typing +import zipfile + +from dataclasses import dataclass +from functools import lru_cache +from typing_extensions import Literal + + +DEFAULT_SIZE = 32 +MODE = Literal['r', 'w', 'x', 'a'] + + +@dataclass(order=True) +class PathInZip: + """Hold a path of file within a zip file. + + Args: + path (str): The convention is :. + Let's assume there is a zip file /some/location/foo.zip + and inside of it is a json file located at /data/file1.json, + Then we expect path = "/some/location/foo.zip:/data/file1.json". + """ + + INFO_PATH_SEP = ':' + zip_path: str + file_path: str + + def __init__(self, path: str) -> None: + split_path = path.split(self.INFO_PATH_SEP) + assert len(split_path) == 2 + self.zip_path, self.file_path = split_path + + @classmethod + def from_paths(cls, zip_path: str, file_path: str): + return cls(zip_path + cls.INFO_PATH_SEP + file_path) + + def __str__(self) -> str: + return self.zip_path + self.INFO_PATH_SEP + self.file_path + + +def _open_zip(path: str, mode: MODE = 'r'): + return zipfile.ZipFile(path, mode) + + +_cached_open_zip = lru_cache(DEFAULT_SIZE)(_open_zip) + + +def set_zip_cache_size(max_size: int): + """Sets the maximal LRU caching for zip file opening. + + Args: + max_size (int): the maximal LRU cache. + """ + global _cached_open_zip + _cached_open_zip = lru_cache(max_size)(_open_zip) + + +def open_file_in_zip(path_in_zip: PathInZip, mode: str = 'r') -> typing.IO: + """Opens a file stored inside a zip and returns a file-like object. + + Args: + path_in_zip (PathInZip): A PathInZip object representing the file to return a file-like object of. + mode (str): The mode in which to open the file with. + Returns: + A file-like object for PathInZip. + """ + zf = _cached_open_zip(path_in_zip.zip_path) + return zf.open(path_in_zip.file_path) diff --git a/src/audiocraft/environment.py b/src/audiocraft/environment.py new file mode 100644 index 0000000000000000000000000000000000000000..adc7819305758bb50a9984928bfa7f13eabef5f5 --- /dev/null +++ b/src/audiocraft/environment.py @@ -0,0 +1,176 @@ +# 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. + +""" +Provides cluster and tools configuration across clusters (slurm, dora, utilities). +""" + +import logging +import os +from pathlib import Path +import re +import typing as tp + +import omegaconf + +from .utils.cluster import _guess_cluster_type + + +logger = logging.getLogger(__name__) + + +class AudioCraftEnvironment: + """Environment configuration for teams and clusters. + + AudioCraftEnvironment picks compute cluster settings (slurm, dora) from the current running environment + or declared variable and the loaded team configuration. Additionally, the AudioCraftEnvironment + provides pointers to a reference folder resolved automatically across clusters that is shared across team members, + allowing to share sigs or other files to run jobs. Finally, it provides dataset mappers to automatically + map dataset file paths to new locations across clusters, allowing to use the same manifest of files across cluters. + + The cluster type is identified automatically and base configuration file is read from config/teams.yaml. + Use the following environment variables to specify the cluster, team or configuration: + + AUDIOCRAFT_CLUSTER (optional): Cluster type to enforce. Useful if the cluster type + cannot be inferred automatically. + AUDIOCRAFT_CONFIG (optional): Path to yaml config holding the teams configuration. + If not set, configuration is read from config/teams.yaml. + AUDIOCRAFT_TEAM (optional): Name of the team. Recommended to set to your own team. + Cluster configuration are shared across teams to match compute allocation, + specify your cluster configuration in the configuration file under a key mapping + your team name. + """ + _instance = None + DEFAULT_TEAM = "default" + + def __init__(self) -> None: + """Loads configuration.""" + self.team: str = os.getenv("AUDIOCRAFT_TEAM", self.DEFAULT_TEAM) + cluster_type = _guess_cluster_type() + cluster = os.getenv( + "AUDIOCRAFT_CLUSTER", cluster_type.value + ) + logger.info("Detecting cluster type %s", cluster_type) + + self.cluster: str = cluster + + config_path = os.getenv( + "AUDIOCRAFT_CONFIG", + Path(__file__) + .parent.parent.joinpath("config/teams", self.team) + .with_suffix(".yaml"), + ) + self.config = omegaconf.OmegaConf.load(config_path) + self._dataset_mappers = [] + cluster_config = self._get_cluster_config() + if "dataset_mappers" in cluster_config: + for pattern, repl in cluster_config["dataset_mappers"].items(): + regex = re.compile(pattern) + self._dataset_mappers.append((regex, repl)) + + def _get_cluster_config(self) -> omegaconf.DictConfig: + assert isinstance(self.config, omegaconf.DictConfig) + return self.config[self.cluster] + + @classmethod + def instance(cls): + if cls._instance is None: + cls._instance = cls() + return cls._instance + + @classmethod + def reset(cls): + """Clears the environment and forces a reload on next invocation.""" + cls._instance = None + + @classmethod + def get_team(cls) -> str: + """Gets the selected team as dictated by the AUDIOCRAFT_TEAM env var. + If not defined, defaults to "labs". + """ + return cls.instance().team + + @classmethod + def get_cluster(cls) -> str: + """Gets the detected cluster. + This value can be overridden by the AUDIOCRAFT_CLUSTER env var. + """ + return cls.instance().cluster + + @classmethod + def get_dora_dir(cls) -> Path: + """Gets the path to the dora directory for the current team and cluster. + Value is overridden by the AUDIOCRAFT_DORA_DIR env var. + """ + cluster_config = cls.instance()._get_cluster_config() + dora_dir = os.getenv("AUDIOCRAFT_DORA_DIR", cluster_config["dora_dir"]) + logger.warning(f"Dora directory: {dora_dir}") + return Path(dora_dir) + + @classmethod + def get_reference_dir(cls) -> Path: + """Gets the path to the reference directory for the current team and cluster. + Value is overridden by the AUDIOCRAFT_REFERENCE_DIR env var. + """ + cluster_config = cls.instance()._get_cluster_config() + return Path(os.getenv("AUDIOCRAFT_REFERENCE_DIR", cluster_config["reference_dir"])) + + @classmethod + def get_slurm_exclude(cls) -> tp.Optional[str]: + """Get the list of nodes to exclude for that cluster.""" + cluster_config = cls.instance()._get_cluster_config() + return cluster_config.get("slurm_exclude") + + @classmethod + def get_slurm_partitions(cls, partition_types: tp.Optional[tp.List[str]] = None) -> str: + """Gets the requested partitions for the current team and cluster as a comma-separated string. + + Args: + partition_types (list[str], optional): partition types to retrieve. Values must be + from ['global', 'team']. If not provided, the global partition is returned. + """ + if not partition_types: + partition_types = ["global"] + + cluster_config = cls.instance()._get_cluster_config() + partitions = [ + cluster_config["partitions"][partition_type] + for partition_type in partition_types + ] + return ",".join(partitions) + + @classmethod + def resolve_reference_path(cls, path: tp.Union[str, Path]) -> Path: + """Converts reference placeholder in path with configured reference dir to resolve paths. + + Args: + path (str or Path): Path to resolve. + Returns: + Path: Resolved path. + """ + path = str(path) + + if path.startswith("//reference"): + reference_dir = cls.get_reference_dir() + logger.warn(f"Reference directory: {reference_dir}") + assert ( + reference_dir.exists() and reference_dir.is_dir() + ), f"Reference directory does not exist: {reference_dir}." + path = re.sub("^//reference", str(reference_dir), path) + + return Path(path) + + @classmethod + def apply_dataset_mappers(cls, path: str) -> str: + """Applies dataset mapping regex rules as defined in the configuration. + If no rules are defined, the path is returned as-is. + """ + instance = cls.instance() + + for pattern, repl in instance._dataset_mappers: + path = pattern.sub(repl, path) + + return path diff --git a/src/audiocraft/models/__init__.py b/src/audiocraft/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d77c01a443819603fb44705b3363dddc6178e97d --- /dev/null +++ b/src/audiocraft/models/__init__.py @@ -0,0 +1,19 @@ +# 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. +""" +Models for EnCodec, AudioGen, MusicGen, as well as the generic LMModel. +""" +# flake8: noqa +# dropped non-inference classes (AudioGen, MultiBandDiffusion, DiffusionUnet/unet.py, +# encodec.py's DAC) to avoid unnecessary dependencies +from . import builders, loaders +from .encodec import ( + CompressionModel, EncodecModel, + HFEncodecModel, HFEncodecCompressionModel) +from .lm import LMModel +from .lm_magnet import MagnetLMModel +from .musicgen import MusicGen +from .magnet import MAGNeT diff --git a/src/audiocraft/models/builders.py b/src/audiocraft/models/builders.py new file mode 100644 index 0000000000000000000000000000000000000000..9dac724ef838d2b6aad3f71fd00c3f43a8ab0a3c --- /dev/null +++ b/src/audiocraft/models/builders.py @@ -0,0 +1,261 @@ +# 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. + +""" +All the functions to build the relevant models and modules +from the Hydra config. +""" + +import typing as tp + +import src.audiocraft as audiocraft +import omegaconf +import torch + +from .encodec import CompressionModel, EncodecModel, InterleaveStereoCompressionModel +from .lm import LMModel +from .lm_magnet import MagnetLMModel +from ..modules.codebooks_patterns import ( + CodebooksPatternProvider, + DelayedPatternProvider, + MusicLMPattern, + ParallelPatternProvider, + UnrolledPatternProvider, + CoarseFirstPattern, +) +from ..modules.conditioners import ( + BaseConditioner, + ChromaStemConditioner, + CLAPEmbeddingConditioner, + ConditionFuser, + ConditioningProvider, + LUTConditioner, + T5Conditioner, +) +from .unet import DiffusionUnet +from .. import quantization as qt +from ..utils.utils import dict_from_config +from ..modules.diffusion_schedule import MultiBandProcessor, SampleProcessor + + +def get_quantizer(quantizer: str, cfg: omegaconf.DictConfig, dimension: int) -> qt.BaseQuantizer: + klass = { + 'no_quant': qt.DummyQuantizer, + 'rvq': qt.ResidualVectorQuantizer + }[quantizer] + kwargs = dict_from_config(getattr(cfg, quantizer)) + if quantizer != 'no_quant': + kwargs['dimension'] = dimension + return klass(**kwargs) + + +def get_encodec_autoencoder(encoder_name: str, cfg: omegaconf.DictConfig): + if encoder_name == 'seanet': + kwargs = dict_from_config(getattr(cfg, 'seanet')) + encoder_override_kwargs = kwargs.pop('encoder') + decoder_override_kwargs = kwargs.pop('decoder') + encoder_kwargs = {**kwargs, **encoder_override_kwargs} + decoder_kwargs = {**kwargs, **decoder_override_kwargs} + encoder = audiocraft.modules.SEANetEncoder(**encoder_kwargs) + decoder = audiocraft.modules.SEANetDecoder(**decoder_kwargs) + return encoder, decoder + else: + raise KeyError(f"Unexpected compression model {cfg.compression_model}") + + +def get_compression_model(cfg: omegaconf.DictConfig) -> CompressionModel: + """Instantiate a compression model.""" + if cfg.compression_model == 'encodec': + kwargs = dict_from_config(getattr(cfg, 'encodec')) + encoder_name = kwargs.pop('autoencoder') + quantizer_name = kwargs.pop('quantizer') + encoder, decoder = get_encodec_autoencoder(encoder_name, cfg) + quantizer = get_quantizer(quantizer_name, cfg, encoder.dimension) + frame_rate = kwargs['sample_rate'] // encoder.hop_length + renormalize = kwargs.pop('renormalize', False) + # deprecated params + kwargs.pop('renorm', None) + return EncodecModel(encoder, decoder, quantizer, + frame_rate=frame_rate, renormalize=renormalize, **kwargs).to(cfg.device) + else: + raise KeyError(f"Unexpected compression model {cfg.compression_model}") + + +def get_lm_model(cfg: omegaconf.DictConfig) -> LMModel: + """Instantiate a transformer LM.""" + if cfg.lm_model in ['transformer_lm', 'transformer_lm_magnet']: + kwargs = dict_from_config(getattr(cfg, 'transformer_lm')) + n_q = kwargs['n_q'] + q_modeling = kwargs.pop('q_modeling', None) + codebooks_pattern_cfg = getattr(cfg, 'codebooks_pattern') + attribute_dropout = dict_from_config(getattr(cfg, 'attribute_dropout')) + cls_free_guidance = dict_from_config(getattr(cfg, 'classifier_free_guidance')) + cfg_prob, cfg_coef = cls_free_guidance['training_dropout'], cls_free_guidance['inference_coef'] + fuser = get_condition_fuser(cfg) + condition_provider = get_conditioner_provider(kwargs["dim"], cfg).to(cfg.device) + if len(fuser.fuse2cond['cross']) > 0: # enforce cross-att programmatically + kwargs['cross_attention'] = True + if codebooks_pattern_cfg.modeling is None: + assert q_modeling is not None, \ + "LM model should either have a codebook pattern defined or transformer_lm.q_modeling" + codebooks_pattern_cfg = omegaconf.OmegaConf.create( + {'modeling': q_modeling, 'delay': {'delays': list(range(n_q))}} + ) + + pattern_provider = get_codebooks_pattern_provider(n_q, codebooks_pattern_cfg) + lm_class = MagnetLMModel if cfg.lm_model == 'transformer_lm_magnet' else LMModel + return lm_class( + pattern_provider=pattern_provider, + condition_provider=condition_provider, + fuser=fuser, + cfg_dropout=cfg_prob, + cfg_coef=cfg_coef, + attribute_dropout=attribute_dropout, + dtype=getattr(torch, cfg.dtype), + device=cfg.device, + **kwargs + ).to(cfg.device) + else: + raise KeyError(f"Unexpected LM model {cfg.lm_model}") + + +def get_conditioner_provider(output_dim: int, cfg: omegaconf.DictConfig) -> ConditioningProvider: + """Instantiate a conditioning model.""" + device = cfg.device + duration = cfg.dataset.segment_duration + cfg = getattr(cfg, 'conditioners') + dict_cfg = {} if cfg is None else dict_from_config(cfg) + conditioners: tp.Dict[str, BaseConditioner] = {} + condition_provider_args = dict_cfg.pop('args', {}) + condition_provider_args.pop('merge_text_conditions_p', None) + condition_provider_args.pop('drop_desc_p', None) + + for cond, cond_cfg in dict_cfg.items(): + model_type = cond_cfg['model'] + model_args = cond_cfg[model_type] + if model_type == 't5': + conditioners[str(cond)] = T5Conditioner(output_dim=output_dim, device=device, **model_args) + elif model_type == 'lut': + conditioners[str(cond)] = LUTConditioner(output_dim=output_dim, **model_args) + elif model_type == 'chroma_stem': + conditioners[str(cond)] = ChromaStemConditioner( + output_dim=output_dim, + duration=duration, + device=device, + **model_args + ) + elif model_type == 'clap': + conditioners[str(cond)] = CLAPEmbeddingConditioner( + output_dim=output_dim, + device=device, + **model_args + ) + else: + raise ValueError(f"Unrecognized conditioning model: {model_type}") + conditioner = ConditioningProvider(conditioners, device=device, **condition_provider_args) + return conditioner + + +def get_condition_fuser(cfg: omegaconf.DictConfig) -> ConditionFuser: + """Instantiate a condition fuser object.""" + fuser_cfg = getattr(cfg, 'fuser') + fuser_methods = ['sum', 'cross', 'prepend', 'input_interpolate'] + fuse2cond = {k: fuser_cfg[k] for k in fuser_methods} + kwargs = {k: v for k, v in fuser_cfg.items() if k not in fuser_methods} + fuser = ConditionFuser(fuse2cond=fuse2cond, **kwargs) + return fuser + + +def get_codebooks_pattern_provider(n_q: int, cfg: omegaconf.DictConfig) -> CodebooksPatternProvider: + """Instantiate a codebooks pattern provider object.""" + pattern_providers = { + 'parallel': ParallelPatternProvider, + 'delay': DelayedPatternProvider, + 'unroll': UnrolledPatternProvider, + 'coarse_first': CoarseFirstPattern, + 'musiclm': MusicLMPattern, + } + name = cfg.modeling + kwargs = dict_from_config(cfg.get(name)) if hasattr(cfg, name) else {} + klass = pattern_providers[name] + return klass(n_q, **kwargs) + + +def get_debug_compression_model(device='cpu', sample_rate: int = 32000): + """Instantiate a debug compression model to be used for unit tests.""" + assert sample_rate in [16000, 32000], "unsupported sample rate for debug compression model" + model_ratios = { + 16000: [10, 8, 8], # 25 Hz at 16kHz + 32000: [10, 8, 16] # 25 Hz at 32kHz + } + ratios: tp.List[int] = model_ratios[sample_rate] + frame_rate = 25 + seanet_kwargs: dict = { + 'n_filters': 4, + 'n_residual_layers': 1, + 'dimension': 32, + 'ratios': ratios, + } + encoder = audiocraft.modules.SEANetEncoder(**seanet_kwargs) + decoder = audiocraft.modules.SEANetDecoder(**seanet_kwargs) + quantizer = qt.ResidualVectorQuantizer(dimension=32, bins=400, n_q=4) + init_x = torch.randn(8, 32, 128) + quantizer(init_x, 1) # initialize kmeans etc. + compression_model = EncodecModel( + encoder, decoder, quantizer, + frame_rate=frame_rate, sample_rate=sample_rate, channels=1).to(device) + return compression_model.eval() + + +def get_diffusion_model(cfg: omegaconf.DictConfig): + # TODO Find a way to infer the channels from dset + channels = cfg.channels + num_steps = cfg.schedule.num_steps + return DiffusionUnet( + chin=channels, num_steps=num_steps, **cfg.diffusion_unet) + + +def get_processor(cfg, sample_rate: int = 24000): + sample_processor = SampleProcessor() + if cfg.use: + kw = dict(cfg) + kw.pop('use') + kw.pop('name') + if cfg.name == "multi_band_processor": + sample_processor = MultiBandProcessor(sample_rate=sample_rate, **kw) + return sample_processor + + +def get_debug_lm_model(device='cpu'): + """Instantiate a debug LM to be used for unit tests.""" + pattern = DelayedPatternProvider(n_q=4) + dim = 16 + providers = { + 'description': LUTConditioner(n_bins=128, dim=dim, output_dim=dim, tokenizer="whitespace"), + } + condition_provider = ConditioningProvider(providers) + fuser = ConditionFuser( + {'cross': ['description'], 'prepend': [], + 'sum': [], 'input_interpolate': []}) + lm = LMModel( + pattern, condition_provider, fuser, + n_q=4, card=400, dim=dim, num_heads=4, custom=True, num_layers=2, + cross_attention=True, causal=True) + return lm.to(device).eval() + + +def get_wrapped_compression_model( + compression_model: CompressionModel, + cfg: omegaconf.DictConfig) -> CompressionModel: + if hasattr(cfg, 'interleave_stereo_codebooks'): + if cfg.interleave_stereo_codebooks.use: + kwargs = dict_from_config(cfg.interleave_stereo_codebooks) + kwargs.pop('use') + compression_model = InterleaveStereoCompressionModel(compression_model, **kwargs) + if hasattr(cfg, 'compression_model_n_q'): + if cfg.compression_model_n_q is not None: + compression_model.set_num_codebooks(cfg.compression_model_n_q) + return compression_model diff --git a/src/audiocraft/models/encodec.py b/src/audiocraft/models/encodec.py new file mode 100644 index 0000000000000000000000000000000000000000..627fddddbb9259ae41fd99cb85693787ec6b1891 --- /dev/null +++ b/src/audiocraft/models/encodec.py @@ -0,0 +1,506 @@ +# 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. +""" + +from abc import ABC, abstractmethod +import logging +import math +from pathlib import Path +import typing as tp + +from einops import rearrange +import numpy as np +import torch +from torch import nn +from transformers import EncodecModel as HFEncodecModel + +from .. import quantization as qt + + +logger = logging.getLogger() + + +class CompressionModel(ABC, nn.Module): + """Base API for all compression models that aim at being used as audio tokenizers + with a language model. + """ + + @abstractmethod + def forward(self, x: torch.Tensor) -> qt.QuantizedResult: + ... + + @abstractmethod + def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]: + """See `EncodecModel.encode`.""" + ... + + @abstractmethod + def decode(self, codes: torch.Tensor, scale: tp.Optional[torch.Tensor] = None): + """See `EncodecModel.decode`.""" + ... + + @abstractmethod + def decode_latent(self, codes: torch.Tensor): + """Decode from the discrete codes to continuous latent space.""" + ... + + @property + @abstractmethod + def channels(self) -> int: + ... + + @property + @abstractmethod + def frame_rate(self) -> float: + ... + + @property + @abstractmethod + def sample_rate(self) -> int: + ... + + @property + @abstractmethod + def cardinality(self) -> int: + ... + + @property + @abstractmethod + def num_codebooks(self) -> int: + ... + + @property + @abstractmethod + def total_codebooks(self) -> int: + ... + + @abstractmethod + def set_num_codebooks(self, n: int): + """Set the active number of codebooks used by the quantizer.""" + ... + + @staticmethod + def get_pretrained( + name: str, device: tp.Union[torch.device, str] = 'cpu' + ) -> 'CompressionModel': + """Instantiate a CompressionModel from a given pretrained model. + + Args: + name (Path or str): name of the pretrained model. See after. + device (torch.device or str): Device on which the model is loaded. + + Pretrained models: + - dac_44khz (https://github.com/descriptinc/descript-audio-codec) + - dac_24khz (same) + - facebook/encodec_24khz (https://huggingface.co/facebook/encodec_24khz) + - facebook/encodec_32khz (https://huggingface.co/facebook/encodec_32khz) + - your own model on Hugging Face. Export instructions to come... + """ + + from . import builders, loaders + model: CompressionModel + if name in ['dac_44khz', 'dac_24khz']: + model_type = name.split('_')[1] + logger.info("Getting pretrained compression model from DAC %s", model_type) + model = DAC(model_type) + elif name in ['debug_compression_model']: + logger.info("Getting pretrained compression model for debug") + model = builders.get_debug_compression_model() + elif Path(name).exists(): + # We assume here if the path exists that it is in fact an AC checkpoint + # that was exported using `audiocraft.utils.export` functions. + model = loaders.load_compression_model(name, device=device) + else: + logger.info("Getting pretrained compression model from HF %s", name) + hf_model = HFEncodecModel.from_pretrained(name) + model = HFEncodecCompressionModel(hf_model).to(device) + return model.to(device).eval() + + +class EncodecModel(CompressionModel): + """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 + sample_rate: int = 0 + channels: int = 0 + + def __init__(self, + encoder: nn.Module, + decoder: nn.Module, + quantizer: qt.BaseQuantizer, + frame_rate: int, + sample_rate: int, + channels: int, + causal: bool = False, + renormalize: bool = False): + super().__init__() + self.encoder = encoder + self.decoder = decoder + self.quantizer = quantizer + self.frame_rate = frame_rate + self.sample_rate = sample_rate + self.channels = channels + self.renormalize = renormalize + self.causal = causal + 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(self, x: torch.Tensor) -> qt.QuantizedResult: + assert x.dim() == 3 + length = x.shape[-1] + x, scale = self.preprocess(x) + + emb = self.encoder(x) + q_res = 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 encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[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 renormalization. + """ + assert x.dim() == 3 + x, scale = self.preprocess(x) + emb = self.encoder(x) + codes = self.quantizer.encode(emb) + return codes, scale + + 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) + + +class DAC(CompressionModel): + def __init__(self, model_type: str = "44khz"): + super().__init__() + try: + import dac.utils + except ImportError: + raise RuntimeError("Could not import dac, make sure it is installed, " + "please run `pip install descript-audio-codec`") + self.model = dac.utils.load_model(model_type=model_type) + self.n_quantizers = self.total_codebooks + self.model.eval() + + def forward(self, x: torch.Tensor) -> qt.QuantizedResult: + # We don't support training with this. + raise NotImplementedError("Forward and training with DAC not supported.") + + def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]: + codes = self.model.encode(x, self.n_quantizers)[1] + return codes[:, :self.n_quantizers], None + + def decode(self, codes: torch.Tensor, scale: tp.Optional[torch.Tensor] = None): + assert scale is None + z_q = self.decode_latent(codes) + return self.model.decode(z_q) + + def decode_latent(self, codes: torch.Tensor): + """Decode from the discrete codes to continuous latent space.""" + return self.model.quantizer.from_codes(codes)[0] + + @property + def channels(self) -> int: + return 1 + + @property + def frame_rate(self) -> float: + return self.model.sample_rate / self.model.hop_length + + @property + def sample_rate(self) -> int: + return self.model.sample_rate + + @property + def cardinality(self) -> int: + return self.model.codebook_size + + @property + def num_codebooks(self) -> int: + return self.n_quantizers + + @property + def total_codebooks(self) -> int: + return self.model.n_codebooks + + def set_num_codebooks(self, n: int): + """Set the active number of codebooks used by the quantizer. + """ + assert n >= 1 + assert n <= self.total_codebooks + self.n_quantizers = n + + +class HFEncodecCompressionModel(CompressionModel): + """Wrapper around HuggingFace Encodec. + """ + def __init__(self, model: HFEncodecModel): + super().__init__() + self.model = model + bws = self.model.config.target_bandwidths + num_codebooks = [ + bw * 1000 / (self.frame_rate * math.log2(self.cardinality)) + for bw in bws + ] + deltas = [nc - int(nc) for nc in num_codebooks] + # Checking we didn't do some bad maths and we indeed have integers! + assert all(deltas) <= 1e-3, deltas + self.possible_num_codebooks = [int(nc) for nc in num_codebooks] + self.set_num_codebooks(max(self.possible_num_codebooks)) + + def forward(self, x: torch.Tensor) -> qt.QuantizedResult: + # We don't support training with this. + raise NotImplementedError("Forward and training with HF EncodecModel not supported.") + + def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]: + bandwidth_index = self.possible_num_codebooks.index(self.num_codebooks) + bandwidth = self.model.config.target_bandwidths[bandwidth_index] + res = self.model.encode(x, None, bandwidth) + assert len(res[0]) == 1 + assert len(res[1]) == 1 + return res[0][0], res[1][0] + + def decode(self, codes: torch.Tensor, scale: tp.Optional[torch.Tensor] = None): + if scale is None: + scales = [None] # type: ignore + else: + scales = scale # type: ignore + res = self.model.decode(codes[None], scales) + return res[0] + + def decode_latent(self, codes: torch.Tensor): + """Decode from the discrete codes to continuous latent space.""" + return self.model.quantizer.decode(codes.transpose(0, 1)) + + @property + def channels(self) -> int: + return self.model.config.audio_channels + + @property + def frame_rate(self) -> float: + hop_length = int(np.prod(self.model.config.upsampling_ratios)) + return self.sample_rate / hop_length + + @property + def sample_rate(self) -> int: + return self.model.config.sampling_rate + + @property + def cardinality(self) -> int: + return self.model.config.codebook_size + + @property + def num_codebooks(self) -> int: + return self._num_codebooks + + @property + def total_codebooks(self) -> int: + return max(self.possible_num_codebooks) + + def set_num_codebooks(self, n: int): + """Set the active number of codebooks used by the quantizer. + """ + if n not in self.possible_num_codebooks: + raise ValueError(f"Allowed values for num codebooks: {self.possible_num_codebooks}") + self._num_codebooks = n + + +class InterleaveStereoCompressionModel(CompressionModel): + """Wraps a CompressionModel to support stereo inputs. The wrapped model + will be applied independently to the left and right channels, and both codebooks + will be interleaved. If the wrapped model returns a representation `[B, K ,T]` per + channel, then the output will be `[B, K * 2, T]` or `[B, K, T * 2]` depending on + `per_timestep`. + + Args: + model (CompressionModel): Compression model to wrap. + per_timestep (bool): Whether to interleave on the timestep dimension + or on the codebooks dimension. + """ + def __init__(self, model: CompressionModel, per_timestep: bool = False): + super().__init__() + self.model = model + self.per_timestep = per_timestep + assert self.model.channels == 1, "Wrapped model is expected to be for monophonic audio" + + @property + def total_codebooks(self): + return self.model.total_codebooks + + @property + def num_codebooks(self): + """Active number of codebooks used by the quantizer. + + ..Warning:: this reports the number of codebooks after the interleaving + of the codebooks! + """ + return self.model.num_codebooks if self.per_timestep else self.model.num_codebooks * 2 + + def set_num_codebooks(self, n: int): + """Set the active number of codebooks used by the quantizer. + + ..Warning:: this sets the number of codebooks before the interleaving! + """ + self.model.set_num_codebooks(n) + + @property + def num_virtual_steps(self) -> float: + """Return the number of virtual steps, e.g. one real step + will be split into that many steps. + """ + return 2 if self.per_timestep else 1 + + @property + def frame_rate(self) -> float: + return self.model.frame_rate * self.num_virtual_steps + + @property + def sample_rate(self) -> int: + return self.model.sample_rate + + @property + def channels(self) -> int: + return 2 + + @property + def cardinality(self): + """Cardinality of each codebook. + """ + return self.model.cardinality + + def forward(self, x: torch.Tensor) -> qt.QuantizedResult: + raise NotImplementedError("Not supported, use encode and decode.") + + def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]: + B, C, T = x.shape + assert C == self.channels, f"Expecting stereo audio but audio num channels is {C}" + + indices_c0, scales_c0 = self.model.encode(x[:, 0, ...].unsqueeze(1)) + indices_c1, scales_c1 = self.model.encode(x[:, 1, ...].unsqueeze(1)) + indices = torch.stack([indices_c0, indices_c1], dim=0) + scales: tp.Optional[torch.Tensor] = None + if scales_c0 is not None and scales_c1 is not None: + scales = torch.stack([scales_c0, scales_c1], dim=1) + + if self.per_timestep: + indices = rearrange(indices, 'c b k t -> b k (t c)', c=2) + else: + indices = rearrange(indices, 'c b k t -> b (k c) t', c=2) + + return (indices, scales) + + def get_left_right_codes(self, codes: torch.Tensor) -> tp.Tuple[torch.Tensor, torch.Tensor]: + if self.per_timestep: + codes = rearrange(codes, 'b k (t c) -> c b k t', c=2) + else: + codes = rearrange(codes, 'b (k c) t -> c b k t', c=2) + return codes[0], codes[1] + + def decode(self, codes: torch.Tensor, scale: tp.Optional[torch.Tensor] = None): + B, K, T = codes.shape + assert T % self.num_virtual_steps == 0, "Provided codes' number of timesteps does not match" + assert K == self.num_codebooks, "Provided codes' number of codebooks does not match" + + scale_c0, scale_c1 = None, None + if scale is not None: + assert scale.size(0) == B and scale.size(1) == 2, f"Scale has unexpected shape: {scale.shape}" + scale_c0 = scale[0, ...] + scale_c1 = scale[1, ...] + + codes_c0, codes_c1 = self.get_left_right_codes(codes) + audio_c0 = self.model.decode(codes_c0, scale_c0) + audio_c1 = self.model.decode(codes_c1, scale_c1) + return torch.cat([audio_c0, audio_c1], dim=1) + + def decode_latent(self, codes: torch.Tensor): + """Decode from the discrete codes to continuous latent space.""" + raise NotImplementedError("Not supported by interleaved stereo wrapped models.") diff --git a/src/audiocraft/models/genmodel.py b/src/audiocraft/models/genmodel.py new file mode 100644 index 0000000000000000000000000000000000000000..96397450b9b03678b35cbf58734ef14b7b33b466 --- /dev/null +++ b/src/audiocraft/models/genmodel.py @@ -0,0 +1,267 @@ +# 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 implementation for audio generative models. This base implementation +combines all the required components to run inference with pretrained audio +generative models. It can be easily inherited by downstream model classes to +provide easy access to the generation API. +""" + +from abc import ABC, abstractmethod +import typing as tp + +import omegaconf +import torch + +from .encodec import CompressionModel +from .lm import LMModel +from .builders import get_wrapped_compression_model +from ..data.audio_utils import convert_audio +from ..modules.conditioners import ConditioningAttributes +from ..utils.autocast import TorchAutocast + + +class BaseGenModel(ABC): + """Base generative model with convenient generation API. + + Args: + name (str): name of the model. + compression_model (CompressionModel): Compression model + used to map audio to invertible discrete representations. + lm (LMModel): Language model over discrete representations. + max_duration (float, optional): maximum duration the model can produce, + otherwise, inferred from the training params. + """ + def __init__(self, name: str, compression_model: CompressionModel, lm: LMModel, + max_duration: tp.Optional[float] = None): + self.name = name + self.compression_model = compression_model + self.lm = lm + self.cfg: tp.Optional[omegaconf.DictConfig] = None + # Just to be safe, let's put everything in eval mode. + self.compression_model.eval() + self.lm.eval() + + if hasattr(lm, 'cfg'): + cfg = lm.cfg + assert isinstance(cfg, omegaconf.DictConfig) + self.cfg = cfg + + if self.cfg is not None: + self.compression_model = get_wrapped_compression_model(self.compression_model, self.cfg) + + if max_duration is None: + if self.cfg is not None: + max_duration = lm.cfg.dataset.segment_duration # type: ignore + else: + raise ValueError("You must provide max_duration when building directly your GenModel") + assert max_duration is not None + + self.max_duration: float = max_duration + self.duration = self.max_duration + + # self.extend_stride is the length of audio extension when generating samples longer + # than self.max_duration. NOTE: the derived class must set self.extend_stride to a + # positive float value when generating with self.duration > self.max_duration. + self.extend_stride: tp.Optional[float] = None + self.device = next(iter(lm.parameters())).device + self.generation_params: dict = {} + self._progress_callback: tp.Optional[tp.Callable[[int, int], None]] = None + if self.device.type == 'cpu': + self.autocast = TorchAutocast(enabled=False) + else: + self.autocast = TorchAutocast( + enabled=True, device_type=self.device.type, dtype=torch.float16) + + @property + def frame_rate(self) -> float: + """Roughly the number of AR steps per seconds.""" + return self.compression_model.frame_rate + + @property + def sample_rate(self) -> int: + """Sample rate of the generated audio.""" + return self.compression_model.sample_rate + + @property + def audio_channels(self) -> int: + """Audio channels of the generated audio.""" + return self.compression_model.channels + + def set_custom_progress_callback(self, progress_callback: tp.Optional[tp.Callable[[int, int], None]] = None): + """Override the default progress callback.""" + self._progress_callback = progress_callback + + @abstractmethod + def set_generation_params(self, *args, **kwargs): + """Set the generation parameters.""" + raise NotImplementedError("No base implementation for setting generation params.") + + @staticmethod + @abstractmethod + def get_pretrained(name: str, device=None): + raise NotImplementedError("No base implementation for getting pretrained model") + + @torch.no_grad() + def _prepare_tokens_and_attributes( + self, + descriptions: tp.Sequence[tp.Optional[str]], + prompt: tp.Optional[torch.Tensor], + ) -> tp.Tuple[tp.List[ConditioningAttributes], tp.Optional[torch.Tensor]]: + """Prepare model inputs. + + Args: + descriptions (list of str): A list of strings used as text conditioning. + prompt (torch.Tensor): A batch of waveforms used for continuation. + """ + attributes = [ + ConditioningAttributes(text={'description': description}) + for description in descriptions] + + if prompt is not None: + if descriptions is not None: + assert len(descriptions) == len(prompt), "Prompt and nb. descriptions doesn't match" + prompt = prompt.to(self.device) + prompt_tokens, scale = self.compression_model.encode(prompt) + assert scale is None + else: + prompt_tokens = None + return attributes, prompt_tokens + + def generate_unconditional(self, num_samples: int, progress: bool = False, + return_tokens: bool = False) -> tp.Union[torch.Tensor, + tp.Tuple[torch.Tensor, torch.Tensor]]: + """Generate samples in an unconditional manner. + + Args: + num_samples (int): Number of samples to be generated. + progress (bool, optional): Flag to display progress of the generation process. Defaults to False. + """ + descriptions: tp.List[tp.Optional[str]] = [None] * num_samples + attributes, prompt_tokens = self._prepare_tokens_and_attributes(descriptions, None) + tokens = self._generate_tokens(attributes, prompt_tokens, progress) + if return_tokens: + return self.generate_audio(tokens), tokens + return self.generate_audio(tokens) + + def generate(self, descriptions: tp.List[str], progress: bool = False, return_tokens: bool = False) \ + -> tp.Union[torch.Tensor, tp.Tuple[torch.Tensor, torch.Tensor]]: + """Generate samples conditioned on text. + + Args: + descriptions (list of str): A list of strings used as text conditioning. + progress (bool, optional): Flag to display progress of the generation process. Defaults to False. + """ + attributes, prompt_tokens = self._prepare_tokens_and_attributes(descriptions, None) + assert prompt_tokens is None + tokens = self._generate_tokens(attributes, prompt_tokens, progress) + if return_tokens: + return self.generate_audio(tokens), tokens + return self.generate_audio(tokens) + + def generate_continuation(self, prompt: torch.Tensor, prompt_sample_rate: int, + descriptions: tp.Optional[tp.List[tp.Optional[str]]] = None, + progress: bool = False, return_tokens: bool = False) \ + -> tp.Union[torch.Tensor, tp.Tuple[torch.Tensor, torch.Tensor]]: + """Generate samples conditioned on audio prompts and an optional text description. + + Args: + prompt (torch.Tensor): A batch of waveforms used for continuation. + Prompt should be [B, C, T], or [C, T] if only one sample is generated. + prompt_sample_rate (int): Sampling rate of the given audio waveforms. + descriptions (list of str, optional): A list of strings used as text conditioning. Defaults to None. + progress (bool, optional): Flag to display progress of the generation process. Defaults to False. + """ + if prompt.dim() == 2: + prompt = prompt[None] + if prompt.dim() != 3: + raise ValueError("prompt should have 3 dimensions: [B, C, T] (C = 1).") + prompt = convert_audio(prompt, prompt_sample_rate, self.sample_rate, self.audio_channels) + if descriptions is None: + descriptions = [None] * len(prompt) + attributes, prompt_tokens = self._prepare_tokens_and_attributes(descriptions, prompt) + assert prompt_tokens is not None + tokens = self._generate_tokens(attributes, prompt_tokens, progress) + if return_tokens: + return self.generate_audio(tokens), tokens + return self.generate_audio(tokens) + + def _generate_tokens(self, attributes: tp.List[ConditioningAttributes], + prompt_tokens: tp.Optional[torch.Tensor], progress: bool = False) -> torch.Tensor: + """Generate discrete audio tokens given audio prompt and/or conditions. + + Args: + attributes (list of ConditioningAttributes): Conditions used for generation (here text). + prompt_tokens (torch.Tensor, optional): Audio prompt used for continuation. + progress (bool, optional): Flag to display progress of the generation process. Defaults to False. + Returns: + torch.Tensor: Generated audio, of shape [B, C, T], T is defined by the generation params. + """ + total_gen_len = int(self.duration * self.frame_rate) + max_prompt_len = int(min(self.duration, self.max_duration) * self.frame_rate) + current_gen_offset: int = 0 + + def _progress_callback(generated_tokens: int, tokens_to_generate: int): + generated_tokens += current_gen_offset + if self._progress_callback is not None: + # Note that total_gen_len might be quite wrong depending on the + # codebook pattern used, but with delay it is almost accurate. + self._progress_callback(generated_tokens, tokens_to_generate) + else: + print(f'{generated_tokens: 6d} / {tokens_to_generate: 6d}', end='\r') + + if prompt_tokens is not None: + assert max_prompt_len >= prompt_tokens.shape[-1], \ + "Prompt is longer than audio to generate" + + callback = None + if progress: + callback = _progress_callback + + if self.duration <= self.max_duration: + # generate by sampling from LM, simple case. + with self.autocast: + gen_tokens = self.lm.generate( + prompt_tokens, attributes, + callback=callback, max_gen_len=total_gen_len, **self.generation_params) + + else: + assert self.extend_stride is not None, "Stride should be defined to generate beyond max_duration" + assert self.extend_stride < self.max_duration, "Cannot stride by more than max generation duration." + all_tokens = [] + if prompt_tokens is None: + prompt_length = 0 + else: + all_tokens.append(prompt_tokens) + prompt_length = prompt_tokens.shape[-1] + + stride_tokens = int(self.frame_rate * self.extend_stride) + while current_gen_offset + prompt_length < total_gen_len: + time_offset = current_gen_offset / self.frame_rate + chunk_duration = min(self.duration - time_offset, self.max_duration) + max_gen_len = int(chunk_duration * self.frame_rate) + with self.autocast: + gen_tokens = self.lm.generate( + prompt_tokens, attributes, + callback=callback, max_gen_len=max_gen_len, **self.generation_params) + if prompt_tokens is None: + all_tokens.append(gen_tokens) + else: + all_tokens.append(gen_tokens[:, :, prompt_tokens.shape[-1]:]) + prompt_tokens = gen_tokens[:, :, stride_tokens:] + prompt_length = prompt_tokens.shape[-1] + current_gen_offset += stride_tokens + + gen_tokens = torch.cat(all_tokens, dim=-1) + return gen_tokens + + def generate_audio(self, gen_tokens: torch.Tensor) -> torch.Tensor: + """Generate Audio from tokens.""" + assert gen_tokens.dim() == 3 + with torch.no_grad(): + gen_audio = self.compression_model.decode(gen_tokens, None) + return gen_audio diff --git a/src/audiocraft/models/lm.py b/src/audiocraft/models/lm.py new file mode 100644 index 0000000000000000000000000000000000000000..3bd6810911e3850369d170421d532d977509c801 --- /dev/null +++ b/src/audiocraft/models/lm.py @@ -0,0 +1,547 @@ +# 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 dataclasses import dataclass +from functools import partial +import logging +import math +import typing as tp + +import torch +from torch import nn + +from ..utils import utils +from ..modules.streaming import StreamingModule, State +from ..modules.transformer import StreamingTransformer, create_norm_fn +from ..modules.conditioners import ( + ConditionFuser, + ClassifierFreeGuidanceDropout, + AttributeDropout, + ConditioningProvider, + ConditioningAttributes, + ConditionType, +) +from ..modules.codebooks_patterns import CodebooksPatternProvider +from ..modules.activations import get_activation_fn + + +logger = logging.getLogger(__name__) +ConditionTensors = tp.Dict[str, ConditionType] +CFGConditions = tp.Union[ConditionTensors, tp.Tuple[ConditionTensors, ConditionTensors]] + + +def get_init_fn(method: str, input_dim: int, init_depth: tp.Optional[int] = None): + """LM layer initialization. + Inspired from xlformers: https://github.com/fairinternal/xlformers + + Args: + method (str): Method name for init function. Valid options are: + 'gaussian', 'uniform'. + input_dim (int): Input dimension of the initialized module. + init_depth (int, optional): Optional init depth value used to rescale + the standard deviation if defined. + """ + # Compute std + std = 1 / math.sqrt(input_dim) + # Rescale with depth + if init_depth is not None: + std = std / math.sqrt(2 * init_depth) + + if method == 'gaussian': + return partial( + torch.nn.init.trunc_normal_, mean=0.0, std=std, a=-3 * std, b=3 * std + ) + elif method == 'uniform': + bound = math.sqrt(3) * std # ensure the standard deviation is `std` + return partial(torch.nn.init.uniform_, a=-bound, b=bound) + else: + raise ValueError("Unsupported layer initialization method") + + +def init_layer(m: nn.Module, + method: str, + init_depth: tp.Optional[int] = None, + zero_bias_init: bool = False): + """Wrapper around ``get_init_fn`` for proper initialization of LM modules. + + Args: + m (nn.Module): Module to initialize. + method (str): Method name for the init function. + init_depth (int, optional): Optional init depth value used to rescale + the standard deviation if defined. + zero_bias_init (bool): Whether to initialize the bias to 0 or not. + """ + if isinstance(m, nn.Linear): + init_fn = get_init_fn(method, m.in_features, init_depth=init_depth) + if m.weight.device.type == 'cpu' and m.weight.dtype == torch.float16: + weight = m.weight.float() + init_fn(weight) + m.weight.data[:] = weight.half() + else: + init_fn(m.weight) + if zero_bias_init and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Embedding): + init_fn = get_init_fn(method, m.embedding_dim, init_depth=None) + if m.weight.device.type == 'cpu' and m.weight.dtype == torch.float16: + weight = m.weight.float() + init_fn(weight) + m.weight.data[:] = weight.half() + else: + init_fn(m.weight) + + +class ScaledEmbedding(nn.Embedding): + """Boost learning rate for embeddings (with `scale`). + """ + def __init__(self, *args, lr=None, **kwargs): + super().__init__(*args, **kwargs) + self.lr = lr + + def make_optim_group(self): + group = {"params": list(self.parameters())} + if self.lr is not None: + group["lr"] = self.lr + return group + + +@dataclass +class LMOutput: + # The logits are already re-aligned with the input codes + # hence no extra shift is required, e.g. when computing CE + logits: torch.Tensor # [B, K, T, card] + mask: torch.Tensor # [B, K, T] + + +class LMModel(StreamingModule): + """Transformer-based language model on multiple streams of codes. + + Args: + pattern_provider (CodebooksPatternProvider): Pattern provider for codebook interleaving. + condition_provider (MusicConditioningProvider): Conditioning provider from metadata. + fuser (ConditionFuser): Fuser handling the fusing of conditions with language model input. + n_q (int): Number of parallel streams to model. + card (int): Cardinality, vocabulary size. + dim (int): Dimension of the transformer encoder. + num_heads (int): Number of heads for the transformer encoder. + hidden_scale (int): Scale for hidden feed forward dimension of the transformer encoder. + norm (str): Normalization method. + norm_first (bool): Use pre-norm instead of post-norm. + emb_lr (float, optional): Embedding-specific learning rate. + bias_proj (bool): Use bias for output projections. + weight_init (str, optional): Method for weight initialization. + depthwise_init (str, optional): Method for depthwise weight initialization. + zero_bias_init (bool): If true and bias in Linears, initialize bias to zeros. + cfg_dropout (float): Classifier-free guidance dropout. + cfg_coef (float): Classifier-free guidance coefficient. + attribute_dropout (dict): Attribute dropout probabilities. + two_step_cfg (bool): Whether to run classifier free-guidance with 2 distinct steps. + **kwargs: Additional parameters for the transformer encoder. + """ + def __init__(self, pattern_provider: CodebooksPatternProvider, condition_provider: ConditioningProvider, + fuser: ConditionFuser, n_q: int = 8, card: int = 1024, dim: int = 128, num_heads: int = 8, + hidden_scale: int = 4, norm: str = 'layer_norm', norm_first: bool = False, + emb_lr: tp.Optional[float] = None, bias_proj: bool = True, + weight_init: tp.Optional[str] = None, depthwise_init: tp.Optional[str] = None, + zero_bias_init: bool = False, cfg_dropout: float = 0, cfg_coef: float = 1.0, + attribute_dropout: tp.Dict[str, tp.Dict[str, float]] = {}, two_step_cfg: bool = False, + **kwargs): + super().__init__() + self.cfg_coef = cfg_coef + self.cfg_dropout = ClassifierFreeGuidanceDropout(p=cfg_dropout) + self.att_dropout = AttributeDropout(p=attribute_dropout) + self.condition_provider = condition_provider + self.fuser = fuser + self.card = card + embed_dim = self.card + 1 + self.n_q = n_q + self.dim = dim + self.pattern_provider = pattern_provider + self.two_step_cfg = two_step_cfg + self.emb = nn.ModuleList([ScaledEmbedding(embed_dim, dim, lr=emb_lr) for _ in range(n_q)]) + if 'activation' in kwargs: + kwargs['activation'] = get_activation_fn(kwargs['activation']) + self.transformer = StreamingTransformer( + d_model=dim, num_heads=num_heads, dim_feedforward=int(hidden_scale * dim), + norm=norm, norm_first=norm_first, **kwargs) + self.out_norm: tp.Optional[nn.Module] = None + if norm_first: + self.out_norm = create_norm_fn(norm, dim) + self.linears = nn.ModuleList([nn.Linear(dim, self.card, bias=bias_proj) for _ in range(n_q)]) + self._init_weights(weight_init, depthwise_init, zero_bias_init) + self._fsdp: tp.Optional[nn.Module] + self.__dict__['_fsdp'] = None + + def _init_weights(self, weight_init: tp.Optional[str], depthwise_init: tp.Optional[str], zero_bias_init: bool): + """Initialization of the transformer module weights. + + Args: + weight_init (str, optional): Weight initialization strategy. See ``get_init_fn`` for valid options. + depthwise_init (str, optional): Depthwise initialization strategy. The following options are valid: + 'current' where the depth corresponds to the current layer index or 'global' where the total number + of layer is used as depth. If not set, no depthwise initialization strategy is used. + zero_bias_init (bool): Whether to initialize bias to zero or not. + """ + assert depthwise_init is None or depthwise_init in ['current', 'global'] + assert depthwise_init is None or weight_init is not None, \ + "If 'depthwise_init' is defined, a 'weight_init' method should be provided." + assert not zero_bias_init or weight_init is not None, \ + "If 'zero_bias_init', a 'weight_init' method should be provided" + + if weight_init is None: + return + + for emb_layer in self.emb: + init_layer(emb_layer, method=weight_init, init_depth=None, zero_bias_init=zero_bias_init) + + for layer_idx, tr_layer in enumerate(self.transformer.layers): + depth = None + if depthwise_init == 'current': + depth = layer_idx + 1 + elif depthwise_init == 'global': + depth = len(self.transformer.layers) + init_fn = partial(init_layer, method=weight_init, init_depth=depth, zero_bias_init=zero_bias_init) + tr_layer.apply(init_fn) + + for linear in self.linears: + init_layer(linear, method=weight_init, init_depth=None, zero_bias_init=zero_bias_init) + + @property + def special_token_id(self) -> int: + return self.card + + @property + def num_codebooks(self) -> int: + return self.n_q + + def forward(self, sequence: torch.Tensor, + conditions: tp.List[ConditioningAttributes], + condition_tensors: tp.Optional[ConditionTensors] = None, + stage: int = -1) -> torch.Tensor: + """Apply language model on sequence and conditions. + Given a tensor of sequence of shape [B, K, S] with K the number of codebooks and + S the sequence steps, return the logits with shape [B, card, K, S]. + + Args: + indices (torch.Tensor): Indices of the codes to model. + conditions (list of ConditioningAttributes): Conditions to use when modeling + the given codes. Note that when evaluating multiple time with the same conditioning + you should pre-compute those and pass them as `condition_tensors`. + condition_tensors (dict[str, ConditionType], optional): Pre-computed conditioning + tensors, see `conditions`. + stage (int): The codebook level that is being predicted. Relevant for MAGNeT + in which prediction is done in a codebook-by-codebook manner. + Takes values in range(n_q), and ignored by default. + Returns: + torch.Tensor: Logits. + """ + B, K, S = sequence.shape + assert K == self.num_codebooks, "Sequence shape must match the specified number of codebooks" + input_ = sum([self.emb[k](sequence[:, k]) for k in range(K)]) + if condition_tensors is None: + assert not self._is_streaming, "Conditions tensors should be precomputed when streaming." + # apply dropout modules + conditions = self.cfg_dropout(conditions) + conditions = self.att_dropout(conditions) + tokenized = self.condition_provider.tokenize(conditions) + # encode conditions and fuse, both have a streaming cache to not recompute when generating. + condition_tensors = self.condition_provider(tokenized) + else: + assert not conditions, "Shouldn't pass both conditions and condition_tensors." + + input_, cross_attention_input = self.fuser(input_, condition_tensors) + + out = self.transformer(input_, cross_attention_src=cross_attention_input, + src_mask=(self.attn_mask_per_stage[stage] if stage >= 0 else None)) + if self.out_norm: + out = self.out_norm(out) + logits = torch.stack([self.linears[k](out) for k in range(K)], dim=1) # [B, K, S, card] + + # remove the prefix from the model outputs + if len(self.fuser.fuse2cond['prepend']) > 0: + logits = logits[:, :, -S:] + + return logits # [B, K, S, card] + + def compute_predictions( + self, codes: torch.Tensor, + conditions: tp.List[ConditioningAttributes], + condition_tensors: tp.Optional[ConditionTensors] = None, + stage: int = -1, + keep_only_valid_steps: bool = True) -> LMOutput: + """Given an input tensor of codes [B, K, T] and list of conditions, runs the model + forward using the specified codes interleaving pattern. + + Args: + codes (torch.Tensor): Input codes of shape [B, K, T] with B the batch size, + K the number of codebooks and T the number of timesteps. + conditions (list of ConditioningAttributes): conditionings to use when modeling + the given codes. Note that when evaluating multiple time with the same conditioning + you should pre-compute those and pass them as `condition_tensors`. + condition_tensors (dict[str, ConditionType], optional): pre-computed conditioning + tensors, see `conditions`. + stage (int): The codebook level that is being predicted. Relevant for MAGNeT + in which prediction is done in a codebook-by-codebook manner. + Takes values in range(n_q), and ignored by default. + 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: + LMOutput: Language model outputs + logits (torch.Tensor) of shape [B, K, T, card] corresponding to the provided codes, + i.e. the first item corresponds to logits to predict the first code, meaning that + no additional shifting of codes and logits is required. + mask (torch.Tensor) of shape [B, K, T], mask over valid and invalid positions. + Given the specified interleaving strategies, parts of the logits and codes should + not be considered as valid predictions because of invalid context. + """ + B, K, T = codes.shape + codes = codes.contiguous() + # map codes [B, K, T] into pattern sequence [B, K, S] using special_token_id for masked tokens + pattern = self.pattern_provider.get_pattern(T) + sequence_codes, sequence_indexes, sequence_mask = pattern.build_pattern_sequence( + codes, self.special_token_id, keep_only_valid_steps=keep_only_valid_steps, + ) + + # apply model on pattern sequence + model = self if self._fsdp is None else self._fsdp + logits = model(sequence_codes, conditions, condition_tensors, stage=stage) # [B, K, S, card] + # map back the logits on pattern sequence to logits on original codes: [B, K, S, card] -> [B, K, T, card] + # and provide the corresponding mask over invalid positions of tokens + logits = logits.permute(0, 3, 1, 2) # [B, card, K, S] + # note: we use nans as special token to make it obvious if we feed unexpected logits + logits, logits_indexes, logits_mask = pattern.revert_pattern_logits( + logits, float('nan'), keep_only_valid_steps=keep_only_valid_steps + ) + logits = logits.permute(0, 2, 3, 1) # [B, K, T, card] + logits_mask = logits_mask[None, :, :].expand(B, -1, -1) # [K, T] -> [B, K, T] + return LMOutput(logits, logits_mask) + + def _sample_next_token(self, + sequence: torch.Tensor, + cfg_conditions: CFGConditions, + unconditional_state: State, + use_sampling: bool = False, + temp: float = 1.0, + top_k: int = 0, + top_p: float = 0.0, + cfg_coef: tp.Optional[float] = None, + two_step_cfg: tp.Optional[bool] = None) -> torch.Tensor: + """Sample next token from the model given a sequence and a set of conditions. The model supports + multiple sampling strategies (greedy sampling, softmax, top-k, top-p...). + + Args: + sequence (torch.Tensor): Current sequence of shape [B, K, S] + with K corresponding to the number of codebooks and S the number of sequence steps. + S = 1 in streaming mode, except for the first step that contains a bigger prompt. + condition_tensors (dict[str, ConditionType): Set of conditions. If CFG is used, + should be twice the batch size, being the concatenation of the conditions + null conditions. + use_sampling (bool): Whether to use a sampling strategy or not. + temp (float): Sampling temperature. + top_k (int): K for "top-k" sampling. + top_p (float): P for "top-p" sampling. + cfg_coef (float, optional): classifier free guidance coefficient + Returns: + next_token (torch.Tensor): Next token tensor of shape [B, K, 1]. + """ + B = sequence.shape[0] + cfg_coef = self.cfg_coef if cfg_coef is None else cfg_coef + model = self if self._fsdp is None else self._fsdp + two_step_cfg = self.two_step_cfg if two_step_cfg is None else two_step_cfg + if two_step_cfg and cfg_conditions != {}: + assert isinstance(cfg_conditions, tuple), type(cfg_conditions) + condition_tensors, null_condition_tensors = cfg_conditions + cond_logits = model(sequence, conditions=[], condition_tensors=condition_tensors) + state = self.get_streaming_state() + self.set_streaming_state(unconditional_state) + uncond_logits = model(sequence, conditions=[], condition_tensors=null_condition_tensors) + unconditional_state.update(self.get_streaming_state()) + self.set_streaming_state(state) + logits = uncond_logits + (cond_logits - uncond_logits) * self.cfg_coef + else: + assert isinstance(cfg_conditions, dict) + condition_tensors = cfg_conditions + if condition_tensors: + # Preparing for CFG, predicting both conditional and unconditional logits. + sequence = torch.cat([sequence, sequence], dim=0) + all_logits = model( + sequence, + conditions=[], condition_tensors=condition_tensors) + if condition_tensors: + cond_logits, uncond_logits = all_logits.split(B, dim=0) # [B, K, T, card] + logits = uncond_logits + (cond_logits - uncond_logits) * cfg_coef + else: + logits = all_logits + + logits = logits.permute(0, 1, 3, 2) # [B, K, card, T] + logits = logits[..., -1] # [B x K x card] + + # Apply softmax for sampling if temp > 0. Else, do greedy sampling to avoid zero division error. + if use_sampling and temp > 0.0: + probs = torch.softmax(logits / temp, dim=-1) + if top_p > 0.0: + next_token = utils.sample_top_p(probs, p=top_p) + elif top_k > 0: + next_token = utils.sample_top_k(probs, k=top_k) + else: + next_token = utils.multinomial(probs, num_samples=1) + else: + next_token = torch.argmax(logits, dim=-1, keepdim=True) + + return next_token + + @torch.no_grad() + def generate(self, + prompt: tp.Optional[torch.Tensor] = None, + conditions: tp.List[ConditioningAttributes] = [], + num_samples: tp.Optional[int] = None, + max_gen_len: int = 256, + use_sampling: bool = True, + temp: float = 1.0, + top_k: int = 250, + top_p: float = 0.0, + cfg_coef: tp.Optional[float] = None, + two_step_cfg: tp.Optional[bool] = None, + remove_prompts: bool = False, + check: bool = False, + callback: tp.Optional[tp.Callable[[int, int], None]] = None, + **kwargs) -> torch.Tensor: + """Generate tokens sampling from the model given a prompt or unconditionally. Generation can + be performed in a greedy fashion or using sampling with top K and top P strategies. + + Args: + prompt (torch.Tensor, optional): Prompt tokens of shape [B, K, T]. + conditions_tensors (list of ConditioningAttributes, optional): List of conditions. + num_samples (int, optional): Number of samples to generate when no prompt and no conditions are given. + max_gen_len (int): Maximum generation length. + use_sampling (bool): Whether to use a sampling strategy or not. + temp (float): Sampling temperature. + top_k (int): K for "top-k" sampling. + top_p (float): P for "top-p" sampling. + cfg_coeff (float, optional): Classifier-free guidance coefficient. + two_step_cfg (bool, optional): Whether to perform classifier-free guidance with two steps generation. + remove_prompts (bool): Whether to remove prompts from generation or not. + check (bool): Whether to apply further checks on generated sequence. + callback (Callback, optional): Callback function to report generation progress. + Returns: + torch.Tensor: Generated tokens. + """ + assert not self.training, "generation shouldn't be used in training mode." + first_param = next(iter(self.parameters())) + device = first_param.device + + # Checking all input shapes are consistent. + possible_num_samples = [] + if num_samples is not None: + possible_num_samples.append(num_samples) + elif prompt is not None: + possible_num_samples.append(prompt.shape[0]) + elif conditions: + possible_num_samples.append(len(conditions)) + else: + possible_num_samples.append(1) + assert [x == possible_num_samples[0] for x in possible_num_samples], "Inconsistent inputs shapes" + num_samples = possible_num_samples[0] + + # below we create set of conditions: one conditional and one unconditional + # to do that we merge the regular condition together with the null condition + # we then do 1 forward pass instead of 2. + # the reason for that is two-fold: + # 1. it is about x2 faster than doing 2 forward passes + # 2. avoid the streaming API treating the 2 passes as part of different time steps + # We also support doing two different passes, in particular to ensure that + # the padding structure is exactly the same between train and test. + # With a batch size of 1, this can be slower though. + cfg_conditions: CFGConditions + two_step_cfg = self.two_step_cfg if two_step_cfg is None else two_step_cfg + if conditions: + null_conditions = ClassifierFreeGuidanceDropout(p=1.0)(conditions) + if two_step_cfg: + cfg_conditions = ( + self.condition_provider(self.condition_provider.tokenize(conditions)), + self.condition_provider(self.condition_provider.tokenize(null_conditions)), + ) + else: + conditions = conditions + null_conditions + tokenized = self.condition_provider.tokenize(conditions) + cfg_conditions = self.condition_provider(tokenized) + else: + cfg_conditions = {} + + if prompt is None: + assert num_samples > 0 + prompt = torch.zeros((num_samples, self.num_codebooks, 0), dtype=torch.long, device=device) + + B, K, T = prompt.shape + start_offset = T + assert start_offset < max_gen_len + + pattern = self.pattern_provider.get_pattern(max_gen_len) + # this token is used as default value for codes that are not generated yet + unknown_token = -1 + + # we generate codes up to the max_gen_len that will be mapped to the pattern sequence + gen_codes = torch.full((B, K, max_gen_len), unknown_token, dtype=torch.long, device=device) + # filling the gen_codes with the prompt if needed + gen_codes[..., :start_offset] = prompt + # create the gen_sequence with proper interleaving from the pattern: [B, K, S] + gen_sequence, indexes, mask = pattern.build_pattern_sequence(gen_codes, self.special_token_id) + # retrieve the start_offset in the sequence: + # it is the first sequence step that contains the `start_offset` timestep + start_offset_sequence = pattern.get_first_step_with_timesteps(start_offset) + assert start_offset_sequence is not None + + with self.streaming(): + unconditional_state = self.get_streaming_state() + prev_offset = 0 + gen_sequence_len = gen_sequence.shape[-1] # gen_sequence shape is [B, K, S] + for offset in range(start_offset_sequence, gen_sequence_len): + # get current sequence (note that the streaming API is providing the caching over previous offsets) + curr_sequence = gen_sequence[..., prev_offset:offset] + curr_mask = mask[None, ..., prev_offset:offset].expand(B, -1, -1) + if check: + # check coherence between mask and sequence + assert (curr_sequence == torch.where(curr_mask, curr_sequence, self.special_token_id)).all() + # should never happen as gen_sequence is filled progressively + assert not (curr_sequence == unknown_token).any() + # sample next token from the model, next token shape is [B, K, 1] + next_token = self._sample_next_token( + curr_sequence, cfg_conditions, unconditional_state, use_sampling, temp, top_k, top_p, + cfg_coef=cfg_coef, two_step_cfg=two_step_cfg) + # ensure the tokens that should be masked are properly set to special_token_id + # as the model never output special_token_id + valid_mask = mask[..., offset:offset+1].expand(B, -1, -1) + next_token[~valid_mask] = self.special_token_id + # ensure we don't overwrite prompt tokens, we only write over unknown tokens + # (then mask tokens should be left as is as well, which is correct) + gen_sequence[..., offset:offset+1] = torch.where( + gen_sequence[..., offset:offset+1] == unknown_token, + next_token, gen_sequence[..., offset:offset+1] + ) + prev_offset = offset + if callback is not None: + callback(1 + offset - start_offset_sequence, gen_sequence_len - start_offset_sequence) + unconditional_state.clear() + + # ensure sequence has been entirely filled + assert not (gen_sequence == unknown_token).any() + # ensure gen_sequence pattern and mask are matching + # which means the gen_sequence is valid according to the pattern + assert ( + gen_sequence == torch.where(mask[None, ...].expand(B, -1, -1), gen_sequence, self.special_token_id) + ).all() + # get back the codes, trimming the prompt if needed and cutting potentially incomplete timesteps + out_codes, out_indexes, out_mask = pattern.revert_pattern_sequence(gen_sequence, special_token=unknown_token) + + # sanity checks over the returned codes and corresponding masks + assert (out_codes[..., :max_gen_len] != unknown_token).all() + assert (out_mask[..., :max_gen_len] == 1).all() + + out_start_offset = start_offset if remove_prompts else 0 + out_codes = out_codes[..., out_start_offset:max_gen_len] + + # ensure the returned codes are all valid + assert (out_codes >= 0).all() and (out_codes <= self.card).all() + return out_codes diff --git a/src/audiocraft/models/lm_magnet.py b/src/audiocraft/models/lm_magnet.py new file mode 100644 index 0000000000000000000000000000000000000000..201a3b76783e0d075d0947902426d4a0cb28eaf7 --- /dev/null +++ b/src/audiocraft/models/lm_magnet.py @@ -0,0 +1,498 @@ +# 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 logging +import math +import typing as tp +import torch +import numpy as np + +from ..utils import utils +from ..modules.conditioners import ( + ClassifierFreeGuidanceDropout, + ConditioningAttributes, + ConditionType, +) +from .lm import LMModel + +logger = logging.getLogger(__name__) +ConditionTensors = tp.Dict[str, ConditionType] +CFGConditions = tp.Union[ConditionTensors, tp.Tuple[ConditionTensors, ConditionTensors]] + + +class MagnetLMModel(LMModel): + """Transformer-based, non-autoregressive model, operates on multiple streams of audio tokens (MAGNeT). + Args: + subcodes_context (int): The number of timesteps attended in the self-attention blocks of codebooks > 0. + When set to -1, attention is unrestricted and all timesteps are attended. Defaults to 5. + compression_model_framerate (int): frame rate of the audio tokenizer. + segment_duration (int): Sample length in seconds. + span_len (int): Determines the length of masking spans. This is the minimal length of consecutive masked tokens, + for both training and inference. Defaults to 3. + **kwargs: Additional parameters for the LMModel. + """ + def __init__(self, subcodes_context: int = 5, compression_model_framerate: int = 50, + segment_duration: int = 10, span_len: int = 3, **kwargs): + super().__init__(**kwargs) + self.causal = kwargs['causal'] + self.subcodes_context = subcodes_context + self.span_len = span_len + self._build_attn_masks(compression_model_framerate=compression_model_framerate, + segment_duration=segment_duration, + num_heads=kwargs['num_heads'], + device=kwargs['device'], dtype=kwargs['dtype']) + + def restricted_context_attn_mask(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + """Creates a restricted attention mask (local attention map) where the context + is determined by self.subcodes_context. + Args: + seq_len (int): token sequence length. + device (torch.device): device of the output tensor. + dtype (torch.dtype): data type of the output tensor. + Returns: + torch.Tensor: The restricted attention mask. + """ + # Return a context restricted non-causal att mask + queries_pos = torch.arange(seq_len, device=device).view(-1, 1) + keys_pos = torch.arange(seq_len, device=device).view(1, -1) + + delta = queries_pos - keys_pos + valid = torch.abs(delta) <= self.subcodes_context + return torch.where( + valid, + torch.zeros([], device=device, dtype=dtype), + torch.full([], float('-inf'), device=device, dtype=dtype)) + + def _stage_attn_mask(self, stage: int, seq_len: int, num_heads: int, + device: torch.device, dtype: torch.dtype) -> tp.Optional[torch.Tensor]: + """Creates a restricted attention mask given the stage (codebook index). + Args: + stage (int): The codebook index. Takes values in [0, n_q]. + seq_len (int): Token sequence length. + num_heads (int): Num transformer attention heads. + device (torch.device): device of the output tensor. + dtype (torch.dtype): data type of the output tensor. + Returns: + torch.Tensor: Either a restricted attention mask or None if stage attention is unrestricted. + """ + sa_mask = None + + if stage > 0 and self.subcodes_context > -1: + # parallel - non-causal - with restricted subcodes context + sa_mask = self.restricted_context_attn_mask(seq_len, device=device, dtype=dtype) + + if sa_mask is not None: + # Repeat for each attention head + sa_mask = sa_mask.repeat((1, num_heads, 1, 1)) + + # align8 to enable memory efficient attention + MEMORY_EFFICIENT_ATTN_ALIGN_FACTOR = 8 + seq_len_aligned = \ + int(np.ceil(seq_len / MEMORY_EFFICIENT_ATTN_ALIGN_FACTOR)) * MEMORY_EFFICIENT_ATTN_ALIGN_FACTOR + + sa_mask_aligned = torch.zeros((1, num_heads, seq_len_aligned, seq_len_aligned), device=device, dtype=dtype) + sa_mask_aligned[..., :seq_len, :seq_len] = sa_mask + sa_mask = sa_mask_aligned + + return sa_mask + + def _build_attn_masks(self, compression_model_framerate: int, segment_duration: int, num_heads: int, + device: torch.device, dtype: torch.dtype): + """Construct attention mask per stage. For each of the RVQ codebook levels in the [0, n_q] range, + either a local attention map or None would be stored as an entry in the self.attn_mask_per_stage list. + Args: + compression_model_framerate (int): The frame rate of the tokenizer. + segment_duration (int): Sample length in seconds. + num_heads (int): Num transformer attention heads. + device (torch.device): device of the output tensor. + dtype (torch.dtype): data type of the output tensor. + """ + seq_len = compression_model_framerate * segment_duration + self.attn_mask_per_stage = [self._stage_attn_mask(stage, seq_len, num_heads, + device, dtype) for stage in range(self.n_q)] + + @torch.no_grad() + def generate(self, + prompt: tp.Optional[torch.Tensor] = None, + conditions: tp.List[ConditioningAttributes] = [], + num_samples: tp.Optional[int] = None, + max_gen_len: int = 256, + use_sampling: bool = True, + temp: float = 1.0, + top_k: int = 250, + top_p: float = 0.0, + cfg_coef: tp.Optional[float] = None, + two_step_cfg: tp.Optional[bool] = None, + remove_prompts: bool = False, + check: bool = False, + callback: tp.Optional[tp.Callable[[int, int], None]] = None, + **kwargs) -> torch.Tensor: + + assert cfg_coef is None, "Unsupported in MAGNeT. Use max_cfg_coef,min_cfg_coef instead." + assert two_step_cfg is None, "MAGNeT currently doesn't support two step classifier-free-guidance." + assert remove_prompts is False, "MAGNeT currently doesn't support the remove_prompts arg." + assert check is False, "MAGNeT currently doesn't support the check arg." + # Call the MAGNeT-specific generation method + return self._generate_magnet(prompt=prompt, + conditions=conditions, + num_samples=num_samples, + max_gen_len=max_gen_len, + use_sampling=use_sampling, + temp=temp, + top_k=top_k, + top_p=top_p, + callback=callback, **kwargs) + + @torch.no_grad() + def _generate_magnet(self, + prompt: tp.Optional[torch.Tensor] = None, + conditions: tp.List[ConditioningAttributes] = [], + num_samples: tp.Optional[int] = None, + max_gen_len: int = 256, + use_sampling: bool = True, + temp: float = 3.0, + top_k: int = 0, + top_p: float = 0.9, + callback: tp.Optional[tp.Callable[[int, int], None]] = None, + max_cfg_coef: float = 10.0, + min_cfg_coef: float = 1.0, + decoding_steps: tp.List[int] = [20, 10, 10, 10], + anneal_temp: bool = True, + span_scoring='max', + span_arrangement='nonoverlap') -> torch.Tensor: + """Generate audio tokens given textual conditions, and optionally given audio prompts, + by running MAGNeT's iterative decoding algorithm for each of the n_q RVQ levels. + Args: + prompt (torch.Tensor): Prompt tokens of shape [B, K, T]. + conditions (list of ConditioningAttributes): List of conditions. + num_samples (int): Number of samples to generate when no prompt and no conditions are given. + max_gen_len (int): Maximum generation length. + use_sampling (bool): Whether to use a sampling strategy or not. + temp (float): Initial sampling temperature. + top_k (int): k for "top-k" sampling. + top_p (float): p for "top-p" sampling. + callback (Callback): Callback function to report generation progress. + max_clsfg_coef (float): Initial coefficient used for classifier free guidance. + min_clsfg_coef (float): Final coefficient used for classifier free guidance. + decoding_steps (list of n_q ints): The number of iterative decoding steps, + for each of the n_q RVQ codebooks. + anneal_temp (bool): When set to True, softmax temperature will be linearly decayed to zero, at each stage. + span_scoring (str): Use the maximum probability of each span ('max') + or the product of probabilities ('prod'). + span_arrangement (str): Use either non-overlapping spans ('nonoverlap') or overlapping spans ('stride1'). + in the masking scheme. + Returns: + torch.Tensor: Generated tokens. + """ + assert not self.training, "generation shouldn't be used in training mode." + first_param = next(iter(self.parameters())) + device = first_param.device + + # Checking all input shapes are consistent. + possible_num_samples = [] + if num_samples is not None: + possible_num_samples.append(num_samples) + elif prompt is not None: + possible_num_samples.append(prompt.shape[0]) + elif conditions: + possible_num_samples.append(len(conditions)) + else: + possible_num_samples.append(1) + assert [x == possible_num_samples[0] for x in possible_num_samples], "Inconsistent inputs shapes" + num_samples = possible_num_samples[0] + + # below we create set of conditions: one conditional and one unconditional + # to do that we merge the regular condition together with the null condition + # we then do 1 forward pass instead of 2. + cfg_conditions: tp.Optional[ConditionTensors] + if conditions: + null_conditions = ClassifierFreeGuidanceDropout(p=1.0)(conditions) + conditions = conditions + null_conditions + tokenized = self.condition_provider.tokenize(conditions) + cfg_conditions = self.condition_provider(tokenized) + else: + cfg_conditions = {} + + if prompt is None: + assert num_samples > 0 + prompt = torch.zeros((num_samples, self.num_codebooks, 0), dtype=torch.long, device=device) + + B, K, prompt_length = prompt.shape + start_offset = prompt_length + assert start_offset < max_gen_len + + mask_id = self.special_token_id + + # we generate codes with a fixed sequence length + shape = (B, K, max_gen_len) + + gen_codes = torch.full(shape, mask_id, dtype=torch.long, device=device) + # filling the gen_codes with the prompt if needed + gen_codes[..., :start_offset] = prompt + # create the gen_sequence with proper interleaving from the pattern: [B, K, S] + gen_sequence = gen_codes + + curr_step = 0 + for stage, n_steps in zip(range(self.n_q), decoding_steps): + gen_sequence, curr_step = self._generate_stage(gen_sequence, + cfg_conditions, + stage=stage, + device=device, + prompt_length=prompt_length, + prompt=prompt, + temp=temp, + max_cfg_coef=max_cfg_coef, + min_cfg_coef=min_cfg_coef, + top_k=top_k, + top_p=top_p, + timesteps=n_steps, + anneal_temp=anneal_temp, + span_scoring=span_scoring, + use_sampling=use_sampling, + span_arrangement=span_arrangement, + curr_step=curr_step, + total_steps=sum(decoding_steps), + callback=callback) + + return gen_sequence + + @torch.no_grad() + def _generate_stage(self, + gen_sequence: torch.Tensor, + condition_tensors: tp.Optional[ConditionTensors], + stage: int, + device: torch.device, + prompt_length: int = 0, + prompt: tp.Optional[torch.Tensor] = None, + use_sampling: bool = True, + temp: float = 3.0, + max_cfg_coef: float = 10.0, + min_cfg_coef: float = 1.0, + top_k: int = 0, + top_p: float = 0.0, + timesteps: int = 10, + anneal_temp: bool = True, + span_scoring: str = 'max', + span_arrangement: str = 'nonoverlap', + curr_step: int = 0, + total_steps: int = 0, + callback: tp.Optional[tp.Callable[[int, int], None]] = None) -> tp.Tuple[torch.Tensor, int]: + """Generate audio tokens of a single RVQ level (stage), given the previously generated stages, + and the textual conditions. + Args: + gen_sequence (torch.Tensor): Previously generated tokens. + condition_tensors (tp.Optional[ConditionTensors]): pre-computed conditioning tensors. + stage (int): RVQ level to generate. + device (torch.device): device of the output tensor. + prompt_length (int): Temporal length of the audio prompt. + prompt (torch.Tensor): Prompt tokens of shape [B, K, T]. + use_sampling (bool): Whether to use a sampling strategy or not. + temp (float): Initial sampling temperature. + max_clsfg_coef (float): Initial coefficient used for classifier free guidance. + min_clsfg_coef (float): Final coefficient used for classifier free guidance. + top_k (int): k for "top-k" sampling. + top_p (float): p for "top-p" sampling. + timesteps (int): Number of iterative decoding steps. + anneal_temp (bool): When set to True, softmax temperature will be linearly decayed to zero, at each stage. + span_scoring (str): Use the maximum probability of each span ('max') + or the product of probabilities ('prod'). + span_arrangement (str): Use either non-overlapping spans ('nonoverlap') or overlapping spans ('stride1'). + in the masking scheme. + curr_step (int): Global iterative decoding step counter. + total_steps (int): Total decoding steps. + callback (Callback): Callback function to report generation progress. + Returns: + tuple(torch.Tensor, int): Generated tokens and the current decoding step counter. + """ + B, K, T = gen_sequence.shape + shape = (B, 1, T) # generating a single codebook per stage + + mask_id = self.special_token_id + stage_gen_seq = torch.full(shape, mask_id, dtype=torch.long, device=device) + + assert span_arrangement == 'nonoverlap' or span_arrangement == 'stride1' + chunk_masking = self.span_len > 1 and span_arrangement == 'nonoverlap' + + DONT_REMASK_ME_SCORE = -1e4 + + model = self if self._fsdp is None else self._fsdp + + if chunk_masking: + # span-wise scores + n_chunks = T // self.span_len + if T % self.span_len != 0: + # trim sequence ending to achieve a multiple of span_len + T = self.span_len * n_chunks + gen_sequence = gen_sequence[..., :T] + stage_gen_seq = stage_gen_seq[..., :T] + + chunked_shape = (B, 1, n_chunks) + n_prompt_chunks = prompt_length // self.span_len + scores = torch.zeros(chunked_shape, dtype=torch.float32, device=device) + scores[..., :n_prompt_chunks] = DONT_REMASK_ME_SCORE + num_chunks_to_gen = n_chunks - n_prompt_chunks + else: + # token-wise scores + scores = torch.zeros(shape, dtype=torch.float32, device=device) + scores[..., :prompt_length] = DONT_REMASK_ME_SCORE + gen_T = T - prompt_length + + # run MAGNeT iterative decoding for "timesteps" iterations + for timestep, steps_left in zip(torch.linspace(0, 1, timesteps, device=device), reversed(range(timesteps))): + + mask_p = torch.cos(timestep * math.pi * 0.5) + + if chunk_masking: + num_masked = max(int((mask_p * num_chunks_to_gen).item()), 1) + else: + num_masked = max(int((mask_p * gen_T).item()), 1) + + # masking + run_lps_masking = (span_arrangement == 'stride1') and self.span_len > 1 + if run_lps_masking: + # masking of the k least probable overlapping (stride 1) spans + mask = torch.concat(( + [self._least_probable_span_masking(scores[[i], :, :], num_masked).to(device) + for i in range(B)]), dim=0) + stage_gen_seq[mask] = mask_id + else: + # masking of the k least probable non-overlapping spans + masked = scores.topk(num_masked, dim=-1).indices + if chunk_masking: + chunks_mask = torch.full(chunked_shape, False, dtype=torch.bool, device=device) + chunks_mask = chunks_mask.scatter(2, masked, True) + mask = torch.repeat_interleave(chunks_mask, self.span_len, dim=-1) + stage_gen_seq[mask] = mask_id + else: + stage_gen_seq = stage_gen_seq.scatter(2, masked, mask_id) + + if prompt is not None: + stage_gen_seq[..., :prompt_length] = prompt[:, stage, :].unsqueeze(1) + + gen_sequence[:, [stage], :] = stage_gen_seq + if condition_tensors: + # duplicate input for classifier free guidance + sequence = torch.cat([gen_sequence, gen_sequence], dim=0) + + all_logits = model(sequence, [], condition_tensors, stage=stage) + + if condition_tensors: + # classifier free guidance with annealing + cond_logits, uncond_logits = all_logits.split(B, dim=0) # [B, K, T, card] + clsfg_coef = float(mask_p) * max_cfg_coef + (1 - float(mask_p)) * min_cfg_coef + logits = uncond_logits + (cond_logits - uncond_logits) * clsfg_coef + else: + logits = all_logits + + # temperature annealing - linear + t = temp * (steps_left / timesteps) if anneal_temp else temp + + # sampling + logits = logits[:, stage, :, :].unsqueeze(1) + probs = torch.softmax(logits / max(t, 1e-2), dim=-1) + if use_sampling: + if top_p > 0.0: + sampled_tokens = utils.sample_top_p(probs, p=top_p) + elif top_k > 0: + sampled_tokens = utils.sample_top_k(probs, k=top_k) + else: + sampled_tokens = utils.multinomial(probs, num_samples=1) + else: + sampled_tokens = torch.argmax(logits, dim=-1, keepdim=True) + + # place mask_id token in each of the masked positions + mask = stage_gen_seq == mask_id + stage_gen_seq = torch.where(mask, sampled_tokens[..., 0], stage_gen_seq) + gen_sequence[:, [stage], :] = stage_gen_seq + + # get probs of sampled tokens + sampled_probs = torch.gather(probs, 3, sampled_tokens)[..., 0] + + # span scoring + if chunk_masking: + if span_scoring == 'max': + # max in linear space + scores = 1 - torch.max(sampled_probs.reshape((B, 1, n_chunks, -1)), dim=-1)[0] + elif span_scoring == 'prod': + # prod in log space + scores = torch.sum(-torch.log(sampled_probs).reshape((B, 1, n_chunks, -1)), dim=-1) + else: + raise NotImplementedError + else: + # prod in log space for lps masking (stride1) + scores = -torch.log(sampled_probs) + + # Fix unmasked tokens by placing inf probs (-inf scores) + if chunk_masking: + scores = scores.masked_fill(~chunks_mask, DONT_REMASK_ME_SCORE) + else: + scores = scores.masked_fill(~mask, DONT_REMASK_ME_SCORE) + + if callback is not None: + curr_step += 1 + callback(curr_step, total_steps) + + return gen_sequence, curr_step + + def _construct_spans_mask(self, span_starts: torch.Tensor, T: int, device: torch.device) -> torch.Tensor: + """Build a [1x1xT] boolean mask consists of overlapping spans of True values, where + span_starts defines the initial index of each span, and the span length is + defined by self.span_len. + Args: + span_starts (torch.Tensor): Boolean mask determines the temporal location of each span start. + T (int): Sequence length. + device (torch.device): device of the output tensor. + Returns: + torch.Tensor: Spans mask of shape [1x1xT] + """ + mask = torch.full((1, 1, T), False, device=device) + mask[:, :, span_starts] = True + shifted_mask = mask.clone() + for _ in range(self.span_len - 1): + shifted_mask = torch.concat((torch.full((1, 1, 1), False, device=device), shifted_mask[:, :, :-1]), dim=-1) + mask = torch.logical_or(mask, shifted_mask) + return mask + + def _least_probable_span_masking(self, scores: torch.Tensor, num_masked_trg: int) -> torch.Tensor: + """Construct a [1x1xT] boolean mask, consists of the u least probable spans, + where the token probability is determined by -scores, and the total + number of masked tokens is as closest as possible to num_masked_trg. + Find u using binary search. + Args: + scores (torch.Tensor): Per token score [-log(prob)] + num_masked_trg: int: The desired amount of tokens to be masked. + Returns: + torch.Tensor: Spans mask of shape [1x1xT] + """ + T = scores.shape[-1] + device = scores.device + scores_unfolded = scores.unfold(2, self.span_len, 1) + # Span score is the product of probs (sum in log space) + span_scores = scores_unfolded.sum(dim=-1) + spans_by_scores = torch.argsort(span_scores[0, 0], descending=True) + + num_masked_trg = max(num_masked_trg, self.span_len) + + # Binary search for u - the number least probable overlapping masked spans s.t. + # the total masking rate is the closest to num_masked_trg / T. + min_u = num_masked_trg // self.span_len + max_u = num_masked_trg - self.span_len + 1 + mid = round(0.5 * (min_u + max_u)) + + if mid == min_u or mid == max_u: + return self._construct_spans_mask(spans_by_scores[:mid], T, device) + + while mid > min_u and mid < max_u: + mask = self._construct_spans_mask(spans_by_scores[:mid], T, device) + n_masked = mask.sum() + if n_masked > num_masked_trg: + max_u = mid + mid = round(0.5 * (min_u + max_u)) + else: + min_u = mid + mid = round(0.5 * (min_u + max_u)) + + return mask diff --git a/src/audiocraft/models/loaders.py b/src/audiocraft/models/loaders.py new file mode 100644 index 0000000000000000000000000000000000000000..5d6392afa102fca73e18613d0078b120762b7f0f --- /dev/null +++ b/src/audiocraft/models/loaders.py @@ -0,0 +1,178 @@ +# 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. + +""" +Utility functions to load from the checkpoints. +Each checkpoint is a torch.saved dict with the following keys: +- 'xp.cfg': the hydra config as dumped during training. This should be used + to rebuild the object using the audiocraft.models.builders functions, +- 'model_best_state': a readily loadable best state for the model, including + the conditioner. The model obtained from `xp.cfg` should be compatible + with this state dict. In the case of a LM, the encodec model would not be + bundled along but instead provided separately. + +Those functions also support loading from a remote location with the Torch Hub API. +They also support overriding some parameters, in particular the device and dtype +of the returned model. +""" + +from pathlib import Path +from huggingface_hub import hf_hub_download +import typing as tp +import os + +from omegaconf import OmegaConf, DictConfig +import torch + +import src.audiocraft as audiocraft +from . import builders +from .encodec import CompressionModel + + +def get_audiocraft_cache_dir() -> tp.Optional[str]: + return os.environ.get('AUDIOCRAFT_CACHE_DIR', None) + + +def _get_state_dict( + file_or_url_or_id: tp.Union[Path, str], + filename: tp.Optional[str] = None, + device='cpu', + cache_dir: tp.Optional[str] = None, +): + if cache_dir is None: + cache_dir = get_audiocraft_cache_dir() + # Return the state dict either from a file or url + file_or_url_or_id = str(file_or_url_or_id) + assert isinstance(file_or_url_or_id, str) + + # weights_only=False: these checkpoints bundle an OmegaConf config, not just tensors + if os.path.isfile(file_or_url_or_id): + return torch.load(file_or_url_or_id, map_location=device, weights_only=False) + + if os.path.isdir(file_or_url_or_id): + file = f"{file_or_url_or_id}/{filename}" + return torch.load(file, map_location=device, weights_only=False) + + elif file_or_url_or_id.startswith('https://'): + return torch.hub.load_state_dict_from_url(file_or_url_or_id, map_location=device, check_hash=True, weights_only=False) + + else: + assert filename is not None, "filename needs to be defined if using HF checkpoints" + + file = hf_hub_download( + repo_id=file_or_url_or_id, filename=filename, cache_dir=cache_dir, + library_name="audiocraft", library_version=audiocraft.__version__) + return torch.load(file, map_location=device, weights_only=False) + + +def load_compression_model_ckpt(file_or_url_or_id: tp.Union[Path, str], cache_dir: tp.Optional[str] = None): + return _get_state_dict(file_or_url_or_id, filename="compression_state_dict.bin", cache_dir=cache_dir) + + +def load_compression_model(file_or_url_or_id: tp.Union[Path, str], device='cpu', cache_dir: tp.Optional[str] = None): + pkg = load_compression_model_ckpt(file_or_url_or_id, cache_dir=cache_dir) + if 'pretrained' in pkg: + return CompressionModel.get_pretrained(pkg['pretrained'], device=device) + cfg = OmegaConf.create(pkg['xp.cfg']) + cfg.device = str(device) + model = builders.get_compression_model(cfg) + model.load_state_dict(pkg['best_state']) + model.eval() + return model + + +def load_lm_model_ckpt(file_or_url_or_id: tp.Union[Path, str], cache_dir: tp.Optional[str] = None): + return _get_state_dict(file_or_url_or_id, filename="state_dict.bin", cache_dir=cache_dir) + + +def _delete_param(cfg: DictConfig, full_name: str): + parts = full_name.split('.') + for part in parts[:-1]: + if part in cfg: + cfg = cfg[part] + else: + return + OmegaConf.set_struct(cfg, False) + if parts[-1] in cfg: + del cfg[parts[-1]] + OmegaConf.set_struct(cfg, True) + + +def load_lm_model(file_or_url_or_id: tp.Union[Path, str], device='cpu', cache_dir: tp.Optional[str] = None): + pkg = load_lm_model_ckpt(file_or_url_or_id, cache_dir=cache_dir) + cfg = OmegaConf.create(pkg['xp.cfg']) + cfg.device = str(device) + if cfg.device == 'cpu': + cfg.dtype = 'float32' + else: + cfg.dtype = 'float16' + _delete_param(cfg, 'conditioners.self_wav.chroma_stem.cache_path') + _delete_param(cfg, 'conditioners.args.merge_text_conditions_p') + _delete_param(cfg, 'conditioners.args.drop_desc_p') + model = builders.get_lm_model(cfg) + model.load_state_dict(pkg['best_state']) + model.eval() + model.cfg = cfg + return model + + +def load_lm_model_magnet(file_or_url_or_id: tp.Union[Path, str], compression_model_frame_rate: int, + device='cpu', cache_dir: tp.Optional[str] = None): + pkg = load_lm_model_ckpt(file_or_url_or_id, cache_dir=cache_dir) + cfg = OmegaConf.create(pkg['xp.cfg']) + cfg.device = str(device) + if cfg.device == 'cpu': + cfg.dtype = 'float32' + else: + cfg.dtype = 'float16' + _delete_param(cfg, 'conditioners.args.merge_text_conditions_p') + _delete_param(cfg, 'conditioners.args.drop_desc_p') + + cfg.transformer_lm.compression_model_framerate = compression_model_frame_rate + cfg.transformer_lm.segment_duration = cfg.dataset.segment_duration + cfg.transformer_lm.span_len = cfg.masking.span_len + + # MAGNeT models v1 support only xformers backend. + from src.audiocraft.modules.transformer import set_efficient_attention_backend + if cfg.transformer_lm.memory_efficient: + set_efficient_attention_backend("xformers") + + model = builders.get_lm_model(cfg) + model.load_state_dict(pkg['best_state']) + model.eval() + model.cfg = cfg + return model + + +def load_mbd_ckpt(file_or_url_or_id: tp.Union[Path, str], + filename: tp.Optional[str] = None, + cache_dir: tp.Optional[str] = None): + return _get_state_dict(file_or_url_or_id, filename=filename, cache_dir=cache_dir) + + +def load_diffusion_models(file_or_url_or_id: tp.Union[Path, str], + device='cpu', + filename: tp.Optional[str] = None, + cache_dir: tp.Optional[str] = None): + pkg = load_mbd_ckpt(file_or_url_or_id, filename=filename, cache_dir=cache_dir) + models = [] + processors = [] + cfgs = [] + sample_rate = pkg['sample_rate'] + for i in range(pkg['n_bands']): + cfg = pkg[i]['cfg'] + model = builders.get_diffusion_model(cfg) + model_dict = pkg[i]['model_state'] + model.load_state_dict(model_dict) + model.to(device) + processor = builders.get_processor(cfg=cfg.processor, sample_rate=sample_rate) + processor_dict = pkg[i]['processor_state'] + processor.load_state_dict(processor_dict) + processor.to(device) + models.append(model) + processors.append(processor) + cfgs.append(cfg) + return models, processors, cfgs diff --git a/src/audiocraft/models/magnet.py b/src/audiocraft/models/magnet.py new file mode 100644 index 0000000000000000000000000000000000000000..453269ad491fa7d05602684f71738df18d0039c1 --- /dev/null +++ b/src/audiocraft/models/magnet.py @@ -0,0 +1,88 @@ +# 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. + +""" +Main model for using MAGNeT. This will combine all the required components +and provide easy access to the generation API. +""" +import typing as tp +import torch + +from .genmodel import BaseGenModel +from .loaders import load_compression_model, load_lm_model_magnet + + +class MAGNeT(BaseGenModel): + """MAGNeT main model with convenient generation API. + Args: + See MusicGen class. + """ + def __init__(self, **kwargs): + super().__init__(**kwargs) + # MAGNeT operates over a fixed sequence length defined in it's config. + self.duration = self.lm.cfg.dataset.segment_duration + self.set_generation_params() + + @staticmethod + def get_pretrained(name: str = 'facebook/magnet-small-10secs', device=None): + """Return pretrained model, we provide six models: + - facebook/magnet-small-10secs (300M), text to music, 10-second audio samples. + # see: https://huggingface.co/facebook/magnet-small-10secs + - facebook/magnet-medium-10secs (1.5B), text to music, 10-second audio samples. + # see: https://huggingface.co/facebook/magnet-medium-10secs + - facebook/magnet-small-30secs (300M), text to music, 30-second audio samples. + # see: https://huggingface.co/facebook/magnet-small-30secs + - facebook/magnet-medium-30secs (1.5B), text to music, 30-second audio samples. + # see: https://huggingface.co/facebook/magnet-medium-30secs + - facebook/audio-magnet-small (300M), text to sound-effect (10-second samples). + # see: https://huggingface.co/facebook/audio-magnet-small + - facebook/audio-magnet-medium (1.5B), text to sound-effect (10-second samples). + # see: https://huggingface.co/facebook/audio-magnet-medium + """ + if device is None: + if torch.cuda.device_count(): + device = 'cuda' + else: + device = 'cpu' + + compression_model = load_compression_model(name, device=device) + lm = load_lm_model_magnet(name, compression_model_frame_rate=int(compression_model.frame_rate), device=device) + + if 'self_wav' in lm.condition_provider.conditioners: + lm.condition_provider.conditioners['self_wav'].match_len_on_eval = True + + kwargs = {'name': name, 'compression_model': compression_model, 'lm': lm} + return MAGNeT(**kwargs) + + def set_generation_params(self, use_sampling: bool = True, top_k: int = 0, + top_p: float = 0.9, temperature: float = 3.0, + max_cfg_coef: float = 10.0, min_cfg_coef: float = 1.0, + decoding_steps: tp.List[int] = [20, 10, 10, 10], + span_arrangement: str = 'nonoverlap'): + """Set the generation parameters for MAGNeT. + + Args: + use_sampling (bool, optional): Use sampling if True, else do argmax decoding. Defaults to True. + top_k (int, optional): top_k used for sampling. Defaults to 0. + top_p (float, optional): top_p used for sampling, when set to 0 top_k is used. Defaults to 0.9. + temperature (float, optional): Initial softmax temperature parameter. Defaults to 3.0. + max_cfg_coef (float, optional): Coefficient used for classifier free guidance. Defaults to 10.0. + min_cfg_coef (float, optional): End coefficient of classifier free guidance annealing. Defaults to 1.0. + decoding_steps (list of n_q ints, optional): The number of iterative decoding steps, + for each of the n_q RVQ codebooks. + span_arrangement (str, optional): Use either non-overlapping spans ('nonoverlap') + or overlapping spans ('stride1') in the masking scheme. + """ + self.generation_params = { + 'use_sampling': use_sampling, + 'temp': temperature, + 'top_k': top_k, + 'top_p': top_p, + 'max_cfg_coef': max_cfg_coef, + 'min_cfg_coef': min_cfg_coef, + 'decoding_steps': [int(s) for s in decoding_steps], + 'span_arrangement': span_arrangement + } diff --git a/src/audiocraft/models/musicgen.py b/src/audiocraft/models/musicgen.py new file mode 100644 index 0000000000000000000000000000000000000000..463f61c6231922b0c8b49d46b81f2c6741743b72 --- /dev/null +++ b/src/audiocraft/models/musicgen.py @@ -0,0 +1,308 @@ +# 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. + +""" +Main model for using MusicGen. This will combine all the required components +and provide easy access to the generation API. +""" + +import typing as tp +import warnings + +import torch + +from .encodec import CompressionModel +from .genmodel import BaseGenModel +from .lm import LMModel +from .builders import get_debug_compression_model, get_debug_lm_model +from .loaders import load_compression_model, load_lm_model +from ..data.audio_utils import convert_audio +from ..modules.conditioners import ConditioningAttributes, WavCondition + + +MelodyList = tp.List[tp.Optional[torch.Tensor]] +MelodyType = tp.Union[torch.Tensor, MelodyList] + + +# backward compatible names mapping +_HF_MODEL_CHECKPOINTS_MAP = { + "small": "facebook/musicgen-small", + "medium": "facebook/musicgen-medium", + "large": "facebook/musicgen-large", + "melody": "facebook/musicgen-melody", +} + + +class MusicGen(BaseGenModel): + """MusicGen main model with convenient generation API. + + Args: + name (str): name of the model. + compression_model (CompressionModel): Compression model + used to map audio to invertible discrete representations. + lm (LMModel): Language model over discrete representations. + max_duration (float, optional): maximum duration the model can produce, + otherwise, inferred from the training params. + """ + def __init__(self, name: str, compression_model: CompressionModel, lm: LMModel, + max_duration: tp.Optional[float] = None): + super().__init__(name, compression_model, lm, max_duration) + self.set_generation_params(duration=15) # default duration + + @staticmethod + def get_pretrained(name: str = 'facebook/musicgen-melody', device=None): + """Return pretrained model, we provide four models: + - facebook/musicgen-small (300M), text to music, + # see: https://huggingface.co/facebook/musicgen-small + - facebook/musicgen-medium (1.5B), text to music, + # see: https://huggingface.co/facebook/musicgen-medium + - facebook/musicgen-melody (1.5B) text to music and text+melody to music, + # see: https://huggingface.co/facebook/musicgen-melody + - facebook/musicgen-large (3.3B), text to music, + # see: https://huggingface.co/facebook/musicgen-large + """ + if device is None: + if torch.cuda.device_count(): + device = 'cuda' + else: + device = 'cpu' + + if name == 'debug': + # used only for unit tests + compression_model = get_debug_compression_model(device) + lm = get_debug_lm_model(device) + return MusicGen(name, compression_model, lm, max_duration=30) + + if name in _HF_MODEL_CHECKPOINTS_MAP: + warnings.warn( + "MusicGen pretrained model relying on deprecated checkpoint mapping. " + + f"Please use full pre-trained id instead: facebook/musicgen-{name}") + name = _HF_MODEL_CHECKPOINTS_MAP[name] + + lm = load_lm_model(name, device=device) + compression_model = load_compression_model(name, device=device) + if 'self_wav' in lm.condition_provider.conditioners: + lm.condition_provider.conditioners['self_wav'].match_len_on_eval = True + lm.condition_provider.conditioners['self_wav']._use_masking = False + + return MusicGen(name, compression_model, lm) + + def set_generation_params(self, use_sampling: bool = True, top_k: int = 250, + top_p: float = 0.0, temperature: float = 1.0, + duration: float = 30.0, cfg_coef: float = 3.0, + two_step_cfg: bool = False, extend_stride: float = 18): + """Set the generation parameters for MusicGen. + + Args: + use_sampling (bool, optional): Use sampling if True, else do argmax decoding. Defaults to True. + top_k (int, optional): top_k used for sampling. Defaults to 250. + top_p (float, optional): top_p used for sampling, when set to 0 top_k is used. Defaults to 0.0. + temperature (float, optional): Softmax temperature parameter. Defaults to 1.0. + duration (float, optional): Duration of the generated waveform. Defaults to 30.0. + cfg_coef (float, optional): Coefficient used for classifier free guidance. Defaults to 3.0. + two_step_cfg (bool, optional): If True, performs 2 forward for Classifier Free Guidance, + instead of batching together the two. This has some impact on how things + are padded but seems to have little impact in practice. + extend_stride: when doing extended generation (i.e. more than 30 seconds), by how much + should we extend the audio each time. Larger values will mean less context is + preserved, and shorter value will require extra computations. + """ + assert extend_stride < self.max_duration, "Cannot stride by more than max generation duration." + self.extend_stride = extend_stride + self.duration = duration + self.generation_params = { + 'use_sampling': use_sampling, + 'temp': temperature, + 'top_k': top_k, + 'top_p': top_p, + 'cfg_coef': cfg_coef, + 'two_step_cfg': two_step_cfg, + } + + def generate_with_chroma(self, descriptions: tp.List[str], melody_wavs: MelodyType, + melody_sample_rate: int, progress: bool = False, + return_tokens: bool = False) -> tp.Union[torch.Tensor, + tp.Tuple[torch.Tensor, torch.Tensor]]: + """Generate samples conditioned on text and melody. + + Args: + descriptions (list of str): A list of strings used as text conditioning. + melody_wavs: (torch.Tensor or list of Tensor): A batch of waveforms used as + melody conditioning. Should have shape [B, C, T] with B matching the description length, + C=1 or 2. It can be [C, T] if there is a single description. It can also be + a list of [C, T] tensors. + melody_sample_rate: (int): Sample rate of the melody waveforms. + progress (bool, optional): Flag to display progress of the generation process. Defaults to False. + """ + if isinstance(melody_wavs, torch.Tensor): + if melody_wavs.dim() == 2: + melody_wavs = melody_wavs[None] + if melody_wavs.dim() != 3: + raise ValueError("Melody wavs should have a shape [B, C, T].") + melody_wavs = list(melody_wavs) + else: + for melody in melody_wavs: + if melody is not None: + assert melody.dim() == 2, "One melody in the list has the wrong number of dims." + + melody_wavs = [ + convert_audio(wav, melody_sample_rate, self.sample_rate, self.audio_channels) + if wav is not None else None + for wav in melody_wavs] + attributes, prompt_tokens = self._prepare_tokens_and_attributes(descriptions=descriptions, prompt=None, + melody_wavs=melody_wavs) + assert prompt_tokens is None + tokens = self._generate_tokens(attributes, prompt_tokens, progress) + if return_tokens: + return self.generate_audio(tokens), tokens + return self.generate_audio(tokens) + + @torch.no_grad() + def _prepare_tokens_and_attributes( + self, + descriptions: tp.Sequence[tp.Optional[str]], + prompt: tp.Optional[torch.Tensor], + melody_wavs: tp.Optional[MelodyList] = None, + ) -> tp.Tuple[tp.List[ConditioningAttributes], tp.Optional[torch.Tensor]]: + """Prepare model inputs. + + Args: + descriptions (list of str): A list of strings used as text conditioning. + prompt (torch.Tensor): A batch of waveforms used for continuation. + melody_wavs (torch.Tensor, optional): A batch of waveforms + used as melody conditioning. Defaults to None. + """ + attributes = [ + ConditioningAttributes(text={'description': description}) + for description in descriptions] + + if melody_wavs is None: + for attr in attributes: + attr.wav['self_wav'] = WavCondition( + torch.zeros((1, 1, 1), device=self.device), + torch.tensor([0], device=self.device), + sample_rate=[self.sample_rate], + path=[None]) + else: + if 'self_wav' not in self.lm.condition_provider.conditioners: + raise RuntimeError("This model doesn't support melody conditioning. " + "Use the `melody` model.") + assert len(melody_wavs) == len(descriptions), \ + f"number of melody wavs must match number of descriptions! " \ + f"got melody len={len(melody_wavs)}, and descriptions len={len(descriptions)}" + for attr, melody in zip(attributes, melody_wavs): + if melody is None: + attr.wav['self_wav'] = WavCondition( + torch.zeros((1, 1, 1), device=self.device), + torch.tensor([0], device=self.device), + sample_rate=[self.sample_rate], + path=[None]) + else: + attr.wav['self_wav'] = WavCondition( + melody[None].to(device=self.device), + torch.tensor([melody.shape[-1]], device=self.device), + sample_rate=[self.sample_rate], + path=[None], + ) + + if prompt is not None: + if descriptions is not None: + assert len(descriptions) == len(prompt), "Prompt and nb. descriptions doesn't match" + prompt = prompt.to(self.device) + prompt_tokens, scale = self.compression_model.encode(prompt) + assert scale is None + else: + prompt_tokens = None + return attributes, prompt_tokens + + def _generate_tokens(self, attributes: tp.List[ConditioningAttributes], + prompt_tokens: tp.Optional[torch.Tensor], progress: bool = False) -> torch.Tensor: + """Generate discrete audio tokens given audio prompt and/or conditions. + + Args: + attributes (list of ConditioningAttributes): Conditions used for generation (text/melody). + prompt_tokens (torch.Tensor, optional): Audio prompt used for continuation. + progress (bool, optional): Flag to display progress of the generation process. Defaults to False. + Returns: + torch.Tensor: Generated audio, of shape [B, C, T], T is defined by the generation params. + """ + total_gen_len = int(self.duration * self.frame_rate) + max_prompt_len = int(min(self.duration, self.max_duration) * self.frame_rate) + current_gen_offset: int = 0 + + def _progress_callback(generated_tokens: int, tokens_to_generate: int): + generated_tokens += current_gen_offset + if self._progress_callback is not None: + # Note that total_gen_len might be quite wrong depending on the + # codebook pattern used, but with delay it is almost accurate. + self._progress_callback(generated_tokens, tokens_to_generate) + else: + print(f'{generated_tokens: 6d} / {tokens_to_generate: 6d}', end='\r') + + if prompt_tokens is not None: + assert max_prompt_len >= prompt_tokens.shape[-1], \ + "Prompt is longer than audio to generate" + + callback = None + if progress: + callback = _progress_callback + + if self.duration <= self.max_duration: + # generate by sampling from LM, simple case. + with self.autocast: + gen_tokens = self.lm.generate( + prompt_tokens, attributes, + callback=callback, max_gen_len=total_gen_len, **self.generation_params) + + else: + # now this gets a bit messier, we need to handle prompts, + # melody conditioning etc. + ref_wavs = [attr.wav['self_wav'] for attr in attributes] + all_tokens = [] + if prompt_tokens is None: + prompt_length = 0 + else: + all_tokens.append(prompt_tokens) + prompt_length = prompt_tokens.shape[-1] + + assert self.extend_stride is not None, "Stride should be defined to generate beyond max_duration" + assert self.extend_stride < self.max_duration, "Cannot stride by more than max generation duration." + stride_tokens = int(self.frame_rate * self.extend_stride) + + while current_gen_offset + prompt_length < total_gen_len: + time_offset = current_gen_offset / self.frame_rate + chunk_duration = min(self.duration - time_offset, self.max_duration) + max_gen_len = int(chunk_duration * self.frame_rate) + for attr, ref_wav in zip(attributes, ref_wavs): + wav_length = ref_wav.length.item() + if wav_length == 0: + continue + # We will extend the wav periodically if it not long enough. + # we have to do it here rather than in conditioners.py as otherwise + # we wouldn't have the full wav. + initial_position = int(time_offset * self.sample_rate) + wav_target_length = int(self.max_duration * self.sample_rate) + positions = torch.arange(initial_position, + initial_position + wav_target_length, device=self.device) + attr.wav['self_wav'] = WavCondition( + ref_wav[0][..., positions % wav_length], + torch.full_like(ref_wav[1], wav_target_length), + [self.sample_rate] * ref_wav[0].size(0), + [None], [0.]) + with self.autocast: + gen_tokens = self.lm.generate( + prompt_tokens, attributes, + callback=callback, max_gen_len=max_gen_len, **self.generation_params) + if prompt_tokens is None: + all_tokens.append(gen_tokens) + else: + all_tokens.append(gen_tokens[:, :, prompt_tokens.shape[-1]:]) + prompt_tokens = gen_tokens[:, :, stride_tokens:] + prompt_length = prompt_tokens.shape[-1] + current_gen_offset += stride_tokens + + gen_tokens = torch.cat(all_tokens, dim=-1) + return gen_tokens diff --git a/src/audiocraft/models/unet.py b/src/audiocraft/models/unet.py new file mode 100644 index 0000000000000000000000000000000000000000..290e68421f0d24a1927e157dfa5ea54f09fe7b56 --- /dev/null +++ b/src/audiocraft/models/unet.py @@ -0,0 +1,214 @@ +# 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. + +""" +Pytorch Unet Module used for diffusion. +""" + +from dataclasses import dataclass +import typing as tp + +import torch +from torch import nn +from torch.nn import functional as F +from src.audiocraft.modules.transformer import StreamingTransformer, create_sin_embedding + + +@dataclass +class Output: + sample: torch.Tensor + + +def get_model(cfg, channels: int, side: int, num_steps: int): + if cfg.model == 'unet': + return DiffusionUnet( + chin=channels, num_steps=num_steps, **cfg.diffusion_unet) + else: + raise RuntimeError('Not Implemented') + + +class ResBlock(nn.Module): + def __init__(self, channels: int, kernel: int = 3, norm_groups: int = 4, + dilation: int = 1, activation: tp.Type[nn.Module] = nn.ReLU, + dropout: float = 0.): + super().__init__() + stride = 1 + padding = dilation * (kernel - stride) // 2 + Conv = nn.Conv1d + Drop = nn.Dropout1d + self.norm1 = nn.GroupNorm(norm_groups, channels) + self.conv1 = Conv(channels, channels, kernel, 1, padding, dilation=dilation) + self.activation1 = activation() + self.dropout1 = Drop(dropout) + + self.norm2 = nn.GroupNorm(norm_groups, channels) + self.conv2 = Conv(channels, channels, kernel, 1, padding, dilation=dilation) + self.activation2 = activation() + self.dropout2 = Drop(dropout) + + def forward(self, x): + h = self.dropout1(self.conv1(self.activation1(self.norm1(x)))) + h = self.dropout2(self.conv2(self.activation2(self.norm2(h)))) + return x + h + + +class DecoderLayer(nn.Module): + def __init__(self, chin: int, chout: int, kernel: int = 4, stride: int = 2, + norm_groups: int = 4, res_blocks: int = 1, activation: tp.Type[nn.Module] = nn.ReLU, + dropout: float = 0.): + super().__init__() + padding = (kernel - stride) // 2 + self.res_blocks = nn.Sequential( + *[ResBlock(chin, norm_groups=norm_groups, dilation=2**idx, dropout=dropout) + for idx in range(res_blocks)]) + self.norm = nn.GroupNorm(norm_groups, chin) + ConvTr = nn.ConvTranspose1d + self.convtr = ConvTr(chin, chout, kernel, stride, padding, bias=False) + self.activation = activation() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.res_blocks(x) + x = self.norm(x) + x = self.activation(x) + x = self.convtr(x) + return x + + +class EncoderLayer(nn.Module): + def __init__(self, chin: int, chout: int, kernel: int = 4, stride: int = 2, + norm_groups: int = 4, res_blocks: int = 1, activation: tp.Type[nn.Module] = nn.ReLU, + dropout: float = 0.): + super().__init__() + padding = (kernel - stride) // 2 + Conv = nn.Conv1d + self.conv = Conv(chin, chout, kernel, stride, padding, bias=False) + self.norm = nn.GroupNorm(norm_groups, chout) + self.activation = activation() + self.res_blocks = nn.Sequential( + *[ResBlock(chout, norm_groups=norm_groups, dilation=2**idx, dropout=dropout) + for idx in range(res_blocks)]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + B, C, T = x.shape + stride, = self.conv.stride + pad = (stride - (T % stride)) % stride + x = F.pad(x, (0, pad)) + + x = self.conv(x) + x = self.norm(x) + x = self.activation(x) + x = self.res_blocks(x) + return x + + +class BLSTM(nn.Module): + """BiLSTM with same hidden units as input dim. + """ + def __init__(self, dim, layers=2): + super().__init__() + self.lstm = nn.LSTM(bidirectional=True, num_layers=layers, hidden_size=dim, input_size=dim) + self.linear = nn.Linear(2 * dim, dim) + + def forward(self, x): + x = x.permute(2, 0, 1) + x = self.lstm(x)[0] + x = self.linear(x) + x = x.permute(1, 2, 0) + return x + + +class DiffusionUnet(nn.Module): + def __init__(self, chin: int = 3, hidden: int = 24, depth: int = 3, growth: float = 2., + max_channels: int = 10_000, num_steps: int = 1000, emb_all_layers=False, cross_attention: bool = False, + bilstm: bool = False, transformer: bool = False, + codec_dim: tp.Optional[int] = None, **kwargs): + super().__init__() + self.encoders = nn.ModuleList() + self.decoders = nn.ModuleList() + self.embeddings: tp.Optional[nn.ModuleList] = None + self.embedding = nn.Embedding(num_steps, hidden) + if emb_all_layers: + self.embeddings = nn.ModuleList() + self.condition_embedding: tp.Optional[nn.Module] = None + for d in range(depth): + encoder = EncoderLayer(chin, hidden, **kwargs) + decoder = DecoderLayer(hidden, chin, **kwargs) + self.encoders.append(encoder) + self.decoders.insert(0, decoder) + if emb_all_layers and d > 0: + assert self.embeddings is not None + self.embeddings.append(nn.Embedding(num_steps, hidden)) + chin = hidden + hidden = min(int(chin * growth), max_channels) + self.bilstm: tp.Optional[nn.Module] + if bilstm: + self.bilstm = BLSTM(chin) + else: + self.bilstm = None + self.use_transformer = transformer + self.cross_attention = False + if transformer: + self.cross_attention = cross_attention + self.transformer = StreamingTransformer(chin, 8, 6, bias_ff=False, bias_attn=False, + cross_attention=cross_attention) + + self.use_codec = False + if codec_dim is not None: + self.conv_codec = nn.Conv1d(codec_dim, chin, 1) + self.use_codec = True + + def forward(self, x: torch.Tensor, step: tp.Union[int, torch.Tensor], condition: tp.Optional[torch.Tensor] = None): + skips = [] + bs = x.size(0) + z = x + view_args = [1] + if type(step) is torch.Tensor: + step_tensor = step + else: + step_tensor = torch.tensor([step], device=x.device, dtype=torch.long).expand(bs) + + for idx, encoder in enumerate(self.encoders): + z = encoder(z) + if idx == 0: + z = z + self.embedding(step_tensor).view(bs, -1, *view_args).expand_as(z) + elif self.embeddings is not None: + z = z + self.embeddings[idx - 1](step_tensor).view(bs, -1, *view_args).expand_as(z) + + skips.append(z) + + if self.use_codec: # insert condition in the bottleneck + assert condition is not None, "Model defined for conditionnal generation" + condition_emb = self.conv_codec(condition) # reshape to the bottleneck dim + assert condition_emb.size(-1) <= 2 * z.size(-1), \ + f"You are downsampling the conditionning with factor >=2 : {condition_emb.size(-1)=} and {z.size(-1)=}" + if not self.cross_attention: + + condition_emb = torch.nn.functional.interpolate(condition_emb, z.size(-1)) + assert z.size() == condition_emb.size() + z += condition_emb + cross_attention_src = None + else: + cross_attention_src = condition_emb.permute(0, 2, 1) # B, T, C + B, T, C = cross_attention_src.shape + positions = torch.arange(T, device=x.device).view(1, -1, 1) + pos_emb = create_sin_embedding(positions, C, max_period=10_000, dtype=cross_attention_src.dtype) + cross_attention_src = cross_attention_src + pos_emb + if self.use_transformer: + z = self.transformer(z.permute(0, 2, 1), cross_attention_src=cross_attention_src).permute(0, 2, 1) + else: + if self.bilstm is None: + z = torch.zeros_like(z) + else: + z = self.bilstm(z) + + for decoder in self.decoders: + s = skips.pop(-1) + z = z[:, :, :s.shape[2]] + z = z + s + z = decoder(z) + + z = z[:, :, :x.shape[2]] + return Output(z) diff --git a/src/audiocraft/modules/__init__.py b/src/audiocraft/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..61418616ef18f0ecca56a007c43af4a731d98b9b --- /dev/null +++ b/src/audiocraft/modules/__init__.py @@ -0,0 +1,22 @@ +# 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. +"""Modules used for building the models.""" + +# flake8: noqa +from .conv import ( + NormConv1d, + NormConv2d, + NormConvTranspose1d, + NormConvTranspose2d, + StreamableConv1d, + StreamableConvTranspose1d, + pad_for_conv1d, + pad1d, + unpad1d, +) +from .lstm import StreamableLSTM +from .seanet import SEANetEncoder, SEANetDecoder +from .transformer import StreamingTransformer \ No newline at end of file diff --git a/src/audiocraft/modules/activations.py b/src/audiocraft/modules/activations.py new file mode 100644 index 0000000000000000000000000000000000000000..2d83d7c4c2dc84c64b724eadbe06157507d4f20d --- /dev/null +++ b/src/audiocraft/modules/activations.py @@ -0,0 +1,96 @@ +# 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 torch +import torch.nn as nn +from torch import Tensor +from typing import Union, Callable + + +class CustomGLU(nn.Module): + """Custom Gated Linear Unit activation. + Applies a modified gated linear unit :math:`a * f(b)` where :math:`a` is the first half + of the input matrices, :math:`b` is the second half, and :math:`f` is a provided activation + function (i.e. sigmoid, swish, etc.). + + Args: + activation (nn.Module): The custom activation to apply in the Gated Linear Unit + dim (int): the dimension on which to split the input. Default: -1 + + Shape: + - Input: :math:`(\ast_1, N, \ast_2)` where `*` means, any number of additional + dimensions + - Output: :math:`(\ast_1, M, \ast_2)` where :math:`M=N/2` + + Examples:: + >>> m = CustomGLU(nn.Sigmoid()) + >>> input = torch.randn(4, 2) + >>> output = m(input) + """ + def __init__(self, activation: nn.Module, dim: int = -1): + super(CustomGLU, self).__init__() + self.dim = dim + self.activation = activation + + def forward(self, x: Tensor): + assert x.shape[self.dim] % 2 == 0 # M = N / 2 + a, b = torch.chunk(x, 2, dim=self.dim) + return a * self.activation(b) + + +class SwiGLU(CustomGLU): + """SiLU Gated Linear Unit activation. + Applies SiLU Gated Linear Unit :math:`a * SiLU(b)` where :math:`a` is + the first half of the input matrices, :math:`b` is the second half. + + Args: + dim (int): the dimension on which to split the input. Default: -1 + """ + def __init__(self, dim: int = -1): + super(SwiGLU, self).__init__(nn.SiLU(), dim) + + +class GeGLU(CustomGLU): + """GeLU Gated Linear Unit activation. + Applies GeLU Gated Linear Unit :math:`a * GELU(b)` where :math:`a` is + the first half of the input matrices, :math:`b` is the second half. + + Args: + dim (int): the dimension on which to split the input. Default: -1 + """ + def __init__(self, dim: int = -1): + super(GeGLU, self).__init__(nn.GELU(), dim) + + +class ReGLU(CustomGLU): + """ReLU Gated Linear Unit activation. + Applies ReLU Gated Linear Unit :math:`a * ReLU(b)` where :math:`a` is + the first half of the input matrices, :math:`b` is the second half. + + Args: + dim (int): the dimension on which to split the input. Default: -1 + """ + def __init__(self, dim: int = -1): + super(ReGLU, self).__init__(nn.ReLU(), dim) + + +def get_activation_fn( + activation: Union[str, Callable[[Tensor], Tensor]] +) -> Union[str, Callable[[Tensor], Tensor]]: + """Helper function to map an activation string to the activation class. + If the supplied activation is not a string that is recognized, the activation is passed back. + + Args: + activation (str, or Callable[[Tensor], Tensor]): Activation to check + """ + if isinstance(activation, str): + if activation == "reglu": + return ReGLU() + elif activation == "geglu": + return GeGLU() + elif activation == "swiglu": + return SwiGLU() + return activation diff --git a/src/audiocraft/modules/chroma.py b/src/audiocraft/modules/chroma.py new file mode 100644 index 0000000000000000000000000000000000000000..e84fb66b4a4aaefb0b3ccac8a9a44c3b20e48f61 --- /dev/null +++ b/src/audiocraft/modules/chroma.py @@ -0,0 +1,66 @@ +# 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 +from librosa import filters +import torch +from torch import nn +import torch.nn.functional as F +import torchaudio + + +class ChromaExtractor(nn.Module): + """Chroma extraction and quantization. + + Args: + sample_rate (int): Sample rate for the chroma extraction. + n_chroma (int): Number of chroma bins for the chroma extraction. + radix2_exp (int): Size of stft window for the chroma extraction (power of 2, e.g. 12 -> 2^12). + nfft (int, optional): Number of FFT. + winlen (int, optional): Window length. + winhop (int, optional): Window hop size. + argmax (bool, optional): Whether to use argmax. Defaults to False. + norm (float, optional): Norm for chroma normalization. Defaults to inf. + """ + def __init__(self, sample_rate: int, n_chroma: int = 12, radix2_exp: int = 12, nfft: tp.Optional[int] = None, + winlen: tp.Optional[int] = None, winhop: tp.Optional[int] = None, argmax: bool = False, + norm: float = torch.inf): + super().__init__() + self.winlen = winlen or 2 ** radix2_exp + self.nfft = nfft or self.winlen + self.winhop = winhop or (self.winlen // 4) + self.sample_rate = sample_rate + self.n_chroma = n_chroma + self.norm = norm + self.argmax = argmax + self.register_buffer('fbanks', torch.from_numpy(filters.chroma(sr=sample_rate, n_fft=self.nfft, tuning=0, + n_chroma=self.n_chroma)), persistent=False) + self.spec = torchaudio.transforms.Spectrogram(n_fft=self.nfft, win_length=self.winlen, + hop_length=self.winhop, power=2, center=True, + pad=0, normalized=True) + + def forward(self, wav: torch.Tensor) -> torch.Tensor: + T = wav.shape[-1] + # in case we are getting a wav that was dropped out (nullified) + # from the conditioner, make sure wav length is no less that nfft + if T < self.nfft: + pad = self.nfft - T + r = 0 if pad % 2 == 0 else 1 + wav = F.pad(wav, (pad // 2, pad // 2 + r), 'constant', 0) + assert wav.shape[-1] == self.nfft, f"expected len {self.nfft} but got {wav.shape[-1]}" + + spec = self.spec(wav).squeeze(1) + raw_chroma = torch.einsum('cf,...ft->...ct', self.fbanks, spec) + norm_chroma = torch.nn.functional.normalize(raw_chroma, p=self.norm, dim=-2, eps=1e-6) + norm_chroma = rearrange(norm_chroma, 'b d t -> b t d') + + if self.argmax: + idx = norm_chroma.argmax(-1, keepdim=True) + norm_chroma[:] = 0 + norm_chroma.scatter_(dim=-1, index=idx, value=1) + + return norm_chroma diff --git a/src/audiocraft/modules/codebooks_patterns.py b/src/audiocraft/modules/codebooks_patterns.py new file mode 100644 index 0000000000000000000000000000000000000000..386df5826937178e29eec670280f8bea57f1a19e --- /dev/null +++ b/src/audiocraft/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) diff --git a/src/audiocraft/modules/conditioners.py b/src/audiocraft/modules/conditioners.py new file mode 100644 index 0000000000000000000000000000000000000000..bfeb9cfa8909b94470321eb4d5bb5dba7738f65a --- /dev/null +++ b/src/audiocraft/modules/conditioners.py @@ -0,0 +1,1416 @@ +# 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 defaultdict +from copy import deepcopy +from dataclasses import dataclass, field +from itertools import chain +import logging +import math +from pathlib import Path +import random +import re +import typing as tp +import warnings + +import einops +from num2words import num2words +import spacy +from transformers import RobertaTokenizer, T5EncoderModel, T5Tokenizer # type: ignore +import torch +from torch import nn +import torch.nn.functional as F +from torch.nn.utils.rnn import pad_sequence + +from .chroma import ChromaExtractor +from .streaming import StreamingModule +from .transformer import create_sin_embedding +from ..data.audio import audio_read +from ..data.audio_dataset import SegmentInfo +from ..data.audio_utils import convert_audio +from ..environment import AudioCraftEnvironment +from ..quantization import ResidualVectorQuantizer +from ..utils.autocast import TorchAutocast +from ..utils.cache import EmbeddingCache +from ..utils.utils import collate, hash_trick, length_to_mask, load_clap_state_dict, warn_once + + +logger = logging.getLogger(__name__) +TextCondition = tp.Optional[str] # a text condition can be a string or None (if doesn't exist) +ConditionType = tp.Tuple[torch.Tensor, torch.Tensor] # condition, mask + + +class WavCondition(tp.NamedTuple): + wav: torch.Tensor + length: torch.Tensor + sample_rate: tp.List[int] + path: tp.List[tp.Optional[str]] = [] + seek_time: tp.List[tp.Optional[float]] = [] + + +class JointEmbedCondition(tp.NamedTuple): + wav: torch.Tensor + text: tp.List[tp.Optional[str]] + length: torch.Tensor + sample_rate: tp.List[int] + path: tp.List[tp.Optional[str]] = [] + seek_time: tp.List[tp.Optional[float]] = [] + + +@dataclass +class ConditioningAttributes: + text: tp.Dict[str, tp.Optional[str]] = field(default_factory=dict) + wav: tp.Dict[str, WavCondition] = field(default_factory=dict) + joint_embed: tp.Dict[str, JointEmbedCondition] = field(default_factory=dict) + + def __getitem__(self, item): + return getattr(self, item) + + @property + def text_attributes(self): + return self.text.keys() + + @property + def wav_attributes(self): + return self.wav.keys() + + @property + def joint_embed_attributes(self): + return self.joint_embed.keys() + + @property + def attributes(self): + return { + "text": self.text_attributes, + "wav": self.wav_attributes, + "joint_embed": self.joint_embed_attributes, + } + + def to_flat_dict(self): + return { + **{f"text.{k}": v for k, v in self.text.items()}, + **{f"wav.{k}": v for k, v in self.wav.items()}, + **{f"joint_embed.{k}": v for k, v in self.joint_embed.items()} + } + + @classmethod + def from_flat_dict(cls, x): + out = cls() + for k, v in x.items(): + kind, att = k.split(".") + out[kind][att] = v + return out + + +class SegmentWithAttributes(SegmentInfo): + """Base class for all dataclasses that are used for conditioning. + All child classes should implement `to_condition_attributes` that converts + the existing attributes to a dataclass of type ConditioningAttributes. + """ + def to_condition_attributes(self) -> ConditioningAttributes: + raise NotImplementedError() + + +def nullify_condition(condition: ConditionType, dim: int = 1): + """Transform an input condition to a null condition. + The way it is done by converting it to a single zero vector similarly + to how it is done inside WhiteSpaceTokenizer and NoopTokenizer. + + Args: + condition (ConditionType): A tuple of condition and mask (tuple[torch.Tensor, torch.Tensor]) + dim (int): The dimension that will be truncated (should be the time dimension) + WARNING!: dim should not be the batch dimension! + Returns: + ConditionType: A tuple of null condition and mask + """ + assert dim != 0, "dim cannot be the batch dimension!" + assert isinstance(condition, tuple) and \ + isinstance(condition[0], torch.Tensor) and \ + isinstance(condition[1], torch.Tensor), "'nullify_condition' got an unexpected input type!" + cond, mask = condition + B = cond.shape[0] + last_dim = cond.dim() - 1 + out = cond.transpose(dim, last_dim) + out = 0. * out[..., :1] + out = out.transpose(dim, last_dim) + mask = torch.zeros((B, 1), device=out.device).int() + assert cond.dim() == out.dim() + return out, mask + + +def nullify_wav(cond: WavCondition) -> WavCondition: + """Transform a WavCondition to a nullified WavCondition. + It replaces the wav by a null tensor, forces its length to 0, and replaces metadata by dummy attributes. + + Args: + cond (WavCondition): Wav condition with wav, tensor of shape [B, T]. + Returns: + WavCondition: Nullified wav condition. + """ + null_wav, _ = nullify_condition((cond.wav, torch.zeros_like(cond.wav)), dim=cond.wav.dim() - 1) + return WavCondition( + wav=null_wav, + length=torch.tensor([0] * cond.wav.shape[0], device=cond.wav.device), + sample_rate=cond.sample_rate, + path=[None] * cond.wav.shape[0], + seek_time=[None] * cond.wav.shape[0], + ) + + +def nullify_joint_embed(embed: JointEmbedCondition) -> JointEmbedCondition: + """Nullify the joint embedding condition by replacing it by a null tensor, forcing its length to 0, + and replacing metadata by dummy attributes. + + Args: + cond (JointEmbedCondition): Joint embedding condition with wav and text, wav tensor of shape [B, C, T]. + """ + null_wav, _ = nullify_condition((embed.wav, torch.zeros_like(embed.wav)), dim=embed.wav.dim() - 1) + return JointEmbedCondition( + wav=null_wav, text=[None] * len(embed.text), + length=torch.LongTensor([0]).to(embed.wav.device), + sample_rate=embed.sample_rate, + path=[None] * embed.wav.shape[0], + seek_time=[0] * embed.wav.shape[0], + ) + + +class Tokenizer: + """Base tokenizer implementation + (in case we want to introduce more advances tokenizers in the future). + """ + def __call__(self, texts: tp.List[tp.Optional[str]]) -> tp.Tuple[torch.Tensor, torch.Tensor]: + raise NotImplementedError() + + +class WhiteSpaceTokenizer(Tokenizer): + """This tokenizer should be used for natural language descriptions. + For example: + ["he didn't, know he's going home.", 'shorter sentence'] => + [[78, 62, 31, 4, 78, 25, 19, 34], + [59, 77, 0, 0, 0, 0, 0, 0]] + """ + PUNCTUATION = "?:!.,;" + + def __init__(self, n_bins: int, pad_idx: int = 0, language: str = "en_core_web_sm", + lemma: bool = True, stopwords: bool = True) -> None: + self.n_bins = n_bins + self.pad_idx = pad_idx + self.lemma = lemma + self.stopwords = stopwords + try: + self.nlp = spacy.load(language) + except IOError: + spacy.cli.download(language) # type: ignore + self.nlp = spacy.load(language) + + @tp.no_type_check + def __call__(self, texts: tp.List[tp.Optional[str]], + return_text: bool = False) -> tp.Tuple[torch.Tensor, torch.Tensor]: + """Take a list of strings and convert them to a tensor of indices. + + Args: + texts (list[str]): List of strings. + return_text (bool, optional): Whether to return text as additional tuple item. Defaults to False. + Returns: + tuple[torch.Tensor, torch.Tensor]: + - Indices of words in the LUT. + - And a mask indicating where the padding tokens are + """ + output, lengths = [], [] + texts = deepcopy(texts) + for i, text in enumerate(texts): + # if current sample doesn't have a certain attribute, replace with pad token + if text is None: + output.append(torch.Tensor([self.pad_idx])) + lengths.append(0) + continue + + # convert numbers to words + text = re.sub(r"(\d+)", lambda x: num2words(int(x.group(0))), text) # type: ignore + # normalize text + text = self.nlp(text) # type: ignore + # remove stopwords + if self.stopwords: + text = [w for w in text if not w.is_stop] # type: ignore + # remove punctuation + text = [w for w in text if w.text not in self.PUNCTUATION] # type: ignore + # lemmatize if needed + text = [getattr(t, "lemma_" if self.lemma else "text") for t in text] # type: ignore + + texts[i] = " ".join(text) + lengths.append(len(text)) + # convert to tensor + tokens = torch.Tensor([hash_trick(w, self.n_bins) for w in text]) + output.append(tokens) + + mask = length_to_mask(torch.IntTensor(lengths)).int() + padded_output = pad_sequence(output, padding_value=self.pad_idx).int().t() + if return_text: + return padded_output, mask, texts # type: ignore + return padded_output, mask + + +class NoopTokenizer(Tokenizer): + """This tokenizer should be used for global conditioners such as: artist, genre, key, etc. + The difference between this and WhiteSpaceTokenizer is that NoopTokenizer does not split + strings, so "Jeff Buckley" will get it's own index. Whereas WhiteSpaceTokenizer will + split it to ["Jeff", "Buckley"] and return an index per word. + + For example: + ["Queen", "ABBA", "Jeff Buckley"] => [43, 55, 101] + ["Metal", "Rock", "Classical"] => [0, 223, 51] + """ + def __init__(self, n_bins: int, pad_idx: int = 0): + self.n_bins = n_bins + self.pad_idx = pad_idx + + def __call__(self, texts: tp.List[tp.Optional[str]]) -> tp.Tuple[torch.Tensor, torch.Tensor]: + output, lengths = [], [] + for text in texts: + # if current sample doesn't have a certain attribute, replace with pad token + if text is None: + output.append(self.pad_idx) + lengths.append(0) + else: + output.append(hash_trick(text, self.n_bins)) + lengths.append(1) + + tokens = torch.LongTensor(output).unsqueeze(1) + mask = length_to_mask(torch.IntTensor(lengths)).int() + return tokens, mask + + +class BaseConditioner(nn.Module): + """Base model for all conditioner modules. + We allow the output dim to be different than the hidden dim for two reasons: + 1) keep our LUTs small when the vocab is large; + 2) make all condition dims consistent. + + Args: + dim (int): Hidden dim of the model. + output_dim (int): Output dim of the conditioner. + """ + def __init__(self, dim: int, output_dim: int): + super().__init__() + self.dim = dim + self.output_dim = output_dim + self.output_proj = nn.Linear(dim, output_dim) + + def tokenize(self, *args, **kwargs) -> tp.Any: + """Should be any part of the processing that will lead to a synchronization + point, e.g. BPE tokenization with transfer to the GPU. + + The returned value will be saved and return later when calling forward(). + """ + raise NotImplementedError() + + def forward(self, inputs: tp.Any) -> ConditionType: + """Gets input that should be used as conditioning (e.g, genre, description or a waveform). + Outputs a ConditionType, after the input data was embedded as a dense vector. + + Returns: + ConditionType: + - A tensor of size [B, T, D] where B is the batch size, T is the length of the + output embedding and D is the dimension of the embedding. + - And a mask indicating where the padding tokens. + """ + raise NotImplementedError() + + +class TextConditioner(BaseConditioner): + ... + + +class LUTConditioner(TextConditioner): + """Lookup table TextConditioner. + + Args: + n_bins (int): Number of bins. + dim (int): Hidden dim of the model (text-encoder/LUT). + output_dim (int): Output dim of the conditioner. + tokenizer (str): Name of the tokenizer. + pad_idx (int, optional): Index for padding token. Defaults to 0. + """ + def __init__(self, n_bins: int, dim: int, output_dim: int, tokenizer: str, pad_idx: int = 0): + super().__init__(dim, output_dim) + self.embed = nn.Embedding(n_bins, dim) + self.tokenizer: Tokenizer + if tokenizer == 'whitespace': + self.tokenizer = WhiteSpaceTokenizer(n_bins, pad_idx=pad_idx) + elif tokenizer == 'noop': + self.tokenizer = NoopTokenizer(n_bins, pad_idx=pad_idx) + else: + raise ValueError(f"unrecognized tokenizer `{tokenizer}`.") + + def tokenize(self, x: tp.List[tp.Optional[str]]) -> tp.Tuple[torch.Tensor, torch.Tensor]: + device = self.embed.weight.device + tokens, mask = self.tokenizer(x) + tokens, mask = tokens.to(device), mask.to(device) + return tokens, mask + + def forward(self, inputs: tp.Tuple[torch.Tensor, torch.Tensor]) -> ConditionType: + tokens, mask = inputs + embeds = self.embed(tokens) + embeds = self.output_proj(embeds) + embeds = (embeds * mask.unsqueeze(-1)) + return embeds, mask + + +class T5Conditioner(TextConditioner): + """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. + device (str): Device for T5 Conditioner. + autocast_dtype (tp.Optional[str], optional): Autocast dtype. + 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, name: str, output_dim: int, finetune: bool, device: str, + autocast_dtype: tp.Optional[str] = 'float32', word_dropout: float = 0., + normalize_text: bool = False): + assert name in self.MODELS, f"Unrecognized t5 model name (should in {self.MODELS})" + super().__init__(self.MODELS_DIMS[name], output_dim) + self.device = device + self.name = name + self.finetune = finetune + self.word_dropout = word_dropout + if autocast_dtype is None or self.device == 'cpu': + self.autocast = TorchAutocast(enabled=False) + if self.device != 'cpu': + logger.warning("T5 has no autocast, this might lead to NaN") + else: + dtype = getattr(torch, autocast_dtype) + assert isinstance(dtype, torch.dtype) + logger.info(f"T5 will be evaluated with autocast as {autocast_dtype}") + self.autocast = TorchAutocast(enabled=True, device_type=self.device, dtype=dtype) + # 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) + 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.to(device) + + self.normalize_text = normalize_text + if normalize_text: + self.text_normalizer = WhiteSpaceTokenizer(1, lemma=True, stopwords=True) + + def tokenize(self, x: tp.List[tp.Optional[str]]) -> tp.Dict[str, torch.Tensor]: + # if current sample doesn't have a certain attribute, replace with empty string + entries: tp.List[str] = [xi if xi is not None else "" for xi in x] + if self.normalize_text: + _, _, entries = self.text_normalizer(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(self.device) + mask = inputs['attention_mask'] + mask[empty_idx, :] = 0 # zero-out index where the input is non-existant + return inputs + + def forward(self, inputs: tp.Dict[str, torch.Tensor]) -> ConditionType: + mask = inputs['attention_mask'] + with torch.set_grad_enabled(self.finetune), self.autocast: + embeds = self.t5(**inputs).last_hidden_state + embeds = self.output_proj(embeds.to(self.output_proj.weight)) + embeds = (embeds * mask.unsqueeze(-1)) + return embeds, mask + + +class WaveformConditioner(BaseConditioner): + """Base class for all conditioners that take a waveform as input. + Classes that inherit must implement `_get_wav_embedding` that outputs + a continuous tensor, and `_downsampling_factor` that returns the down-sampling + factor of the embedding model. + + Args: + dim (int): The internal representation dimension. + output_dim (int): Output dimension. + device (tp.Union[torch.device, str]): Device. + """ + def __init__(self, dim: int, output_dim: int, device: tp.Union[torch.device, str]): + super().__init__(dim, output_dim) + self.device = device + # if False no masking is done, used in ChromaStemConditioner when completing by periodicity a sample. + self._use_masking = True + + def tokenize(self, x: WavCondition) -> WavCondition: + wav, length, sample_rate, path, seek_time = x + assert length is not None + return WavCondition(wav.to(self.device), length.to(self.device), sample_rate, path, seek_time) + + def _get_wav_embedding(self, x: WavCondition) -> torch.Tensor: + """Gets as input a WavCondition and returns a dense embedding.""" + raise NotImplementedError() + + def _downsampling_factor(self): + """Returns the downsampling factor of the embedding model.""" + raise NotImplementedError() + + def forward(self, x: WavCondition) -> ConditionType: + """Extract condition embedding and mask from a waveform and its metadata. + Args: + x (WavCondition): Waveform condition containing raw waveform and metadata. + Returns: + ConditionType: a dense vector representing the conditioning along with its mask + """ + wav, lengths, *_ = x + with torch.no_grad(): + embeds = self._get_wav_embedding(x) + embeds = embeds.to(self.output_proj.weight) + embeds = self.output_proj(embeds) + + if lengths is not None and self._use_masking: + lengths = lengths / self._downsampling_factor() + mask = length_to_mask(lengths, max_len=embeds.shape[1]).int() # type: ignore + else: + mask = torch.ones_like(embeds[..., 0]) + embeds = (embeds * mask.unsqueeze(-1)) + return embeds, mask + + +class ChromaStemConditioner(WaveformConditioner): + """Chroma conditioner based on stems. + The ChromaStemConditioner uses DEMUCS to first filter out drums and bass, as + the drums and bass often dominate the chroma leading to the chroma features + not containing information about the melody. + + Args: + output_dim (int): Output dimension for the conditioner. + sample_rate (int): Sample rate for the chroma extractor. + n_chroma (int): Number of chroma bins for the chroma extractor. + radix2_exp (int): Size of stft window for the chroma extractor (power of 2, e.g. 12 -> 2^12). + duration (int): duration used during training. This is later used for correct padding + in case we are using chroma as prefix. + match_len_on_eval (bool, optional): if True then all chromas are padded to the training + duration. Defaults to False. + eval_wavs (str, optional): path to a dataset manifest with waveform, this waveforms are used as + conditions during eval (for cases where we don't want to leak test conditions like MusicCaps). + Defaults to None. + n_eval_wavs (int, optional): limits the number of waveforms used for conditioning. Defaults to 0. + device (tp.Union[torch.device, str], optional): Device for the conditioner. + **kwargs: Additional parameters for the chroma extractor. + """ + def __init__(self, output_dim: int, sample_rate: int, n_chroma: int, radix2_exp: int, + duration: float, match_len_on_eval: bool = True, eval_wavs: tp.Optional[str] = None, + n_eval_wavs: int = 0, cache_path: tp.Optional[tp.Union[str, Path]] = None, + device: tp.Union[torch.device, str] = 'cpu', **kwargs): + from demucs import pretrained + super().__init__(dim=n_chroma, output_dim=output_dim, device=device) + self.autocast = TorchAutocast(enabled=device != 'cpu', device_type=self.device, dtype=torch.float32) + self.sample_rate = sample_rate + self.match_len_on_eval = match_len_on_eval + if match_len_on_eval: + self._use_masking = False + self.duration = duration + self.__dict__['demucs'] = pretrained.get_model('htdemucs').to(device) + stem_sources: list = self.demucs.sources # type: ignore + self.stem_indices = torch.LongTensor([stem_sources.index('vocals'), stem_sources.index('other')]).to(device) + self.chroma = ChromaExtractor(sample_rate=sample_rate, n_chroma=n_chroma, + radix2_exp=radix2_exp, **kwargs).to(device) + self.chroma_len = self._get_chroma_len() + self.eval_wavs: tp.Optional[torch.Tensor] = self._load_eval_wavs(eval_wavs, n_eval_wavs) + self.cache = None + if cache_path is not None: + self.cache = EmbeddingCache(Path(cache_path) / 'wav', self.device, + compute_embed_fn=self._get_full_chroma_for_cache, + extract_embed_fn=self._extract_chroma_chunk) + + def _downsampling_factor(self) -> int: + return self.chroma.winhop + + def _load_eval_wavs(self, path: tp.Optional[str], num_samples: int) -> tp.Optional[torch.Tensor]: + """Load pre-defined waveforms from a json. + These waveforms will be used for chroma extraction during evaluation. + This is done to make the evaluation on MusicCaps fair (we shouldn't see the chromas of MusicCaps). + """ + if path is None: + return None + + logger.info(f"Loading evaluation wavs from {path}") + from src.audiocraft.data.audio_dataset import AudioDataset + dataset: AudioDataset = AudioDataset.from_meta( + path, segment_duration=self.duration, min_audio_duration=self.duration, + sample_rate=self.sample_rate, channels=1) + + if len(dataset) > 0: + eval_wavs = dataset.collater([dataset[i] for i in range(num_samples)]).to(self.device) + logger.info(f"Using {len(eval_wavs)} evaluation wavs for chroma-stem conditioner") + return eval_wavs + else: + raise ValueError("Could not find evaluation wavs, check lengths of wavs") + + def reset_eval_wavs(self, eval_wavs: tp.Optional[torch.Tensor]) -> None: + self.eval_wavs = eval_wavs + + def has_eval_wavs(self) -> bool: + return self.eval_wavs is not None + + def _sample_eval_wavs(self, num_samples: int) -> torch.Tensor: + """Sample wavs from a predefined list.""" + assert self.eval_wavs is not None, "Cannot sample eval wavs as no eval wavs provided." + total_eval_wavs = len(self.eval_wavs) + out = self.eval_wavs + if num_samples > total_eval_wavs: + out = self.eval_wavs.repeat(num_samples // total_eval_wavs + 1, 1, 1) + return out[torch.randperm(len(out))][:num_samples] + + def _get_chroma_len(self) -> int: + """Get length of chroma during training.""" + dummy_wav = torch.zeros((1, int(self.sample_rate * self.duration)), device=self.device) + dummy_chr = self.chroma(dummy_wav) + return dummy_chr.shape[1] + + @torch.no_grad() + def _get_stemmed_wav(self, wav: torch.Tensor, sample_rate: int) -> torch.Tensor: + """Get parts of the wav that holds the melody, extracting the main stems from the wav.""" + from demucs.apply import apply_model + from demucs.audio import convert_audio + with self.autocast: + wav = convert_audio( + wav, sample_rate, self.demucs.samplerate, self.demucs.audio_channels) # type: ignore + stems = apply_model(self.demucs, wav, device=self.device) + stems = stems[:, self.stem_indices] # extract relevant stems for melody conditioning + mix_wav = stems.sum(1) # merge extracted stems to single waveform + mix_wav = convert_audio(mix_wav, self.demucs.samplerate, self.sample_rate, 1) # type: ignore + return mix_wav + + @torch.no_grad() + def _extract_chroma(self, wav: torch.Tensor) -> torch.Tensor: + """Extract chroma features from the waveform.""" + with self.autocast: + return self.chroma(wav) + + @torch.no_grad() + def _compute_wav_embedding(self, wav: torch.Tensor, sample_rate: int) -> torch.Tensor: + """Compute wav embedding, applying stem and chroma extraction.""" + # avoid 0-size tensors when we are working with null conds + if wav.shape[-1] == 1: + return self._extract_chroma(wav) + stems = self._get_stemmed_wav(wav, sample_rate) + chroma = self._extract_chroma(stems) + return chroma + + @torch.no_grad() + def _get_full_chroma_for_cache(self, path: tp.Union[str, Path], x: WavCondition, idx: int) -> torch.Tensor: + """Extract chroma from the whole audio waveform at the given path.""" + wav, sr = audio_read(path) + wav = wav[None].to(self.device) + wav = convert_audio(wav, sr, self.sample_rate, to_channels=1) + chroma = self._compute_wav_embedding(wav, self.sample_rate)[0] + return chroma + + def _extract_chroma_chunk(self, full_chroma: torch.Tensor, x: WavCondition, idx: int) -> torch.Tensor: + """Extract a chunk of chroma from the full chroma derived from the full waveform.""" + wav_length = x.wav.shape[-1] + seek_time = x.seek_time[idx] + assert seek_time is not None, ( + "WavCondition seek_time is required " + "when extracting chroma chunks from pre-computed chroma.") + full_chroma = full_chroma.float() + frame_rate = self.sample_rate / self._downsampling_factor() + target_length = int(frame_rate * wav_length / self.sample_rate) + index = int(frame_rate * seek_time) + out = full_chroma[index: index + target_length] + out = F.pad(out[None], (0, 0, 0, target_length - out.shape[0]))[0] + return out.to(self.device) + + @torch.no_grad() + def _get_wav_embedding(self, x: WavCondition) -> torch.Tensor: + """Get the wav embedding from the WavCondition. + The conditioner will either extract the embedding on-the-fly computing it from the condition wav directly + or will rely on the embedding cache to load the pre-computed embedding if relevant. + """ + sampled_wav: tp.Optional[torch.Tensor] = None + if not self.training and self.eval_wavs is not None: + warn_once(logger, "Using precomputed evaluation wavs!") + sampled_wav = self._sample_eval_wavs(len(x.wav)) + + no_undefined_paths = all(p is not None for p in x.path) + no_nullified_cond = x.wav.shape[-1] > 1 + if sampled_wav is not None: + chroma = self._compute_wav_embedding(sampled_wav, self.sample_rate) + elif self.cache is not None and no_undefined_paths and no_nullified_cond: + paths = [Path(p) for p in x.path if p is not None] + chroma = self.cache.get_embed_from_cache(paths, x) + else: + assert all(sr == x.sample_rate[0] for sr in x.sample_rate), "All sample rates in batch should be equal." + chroma = self._compute_wav_embedding(x.wav, x.sample_rate[0]) + + if self.match_len_on_eval: + B, T, C = chroma.shape + if T > self.chroma_len: + chroma = chroma[:, :self.chroma_len] + logger.debug(f"Chroma was truncated to match length! ({T} -> {chroma.shape[1]})") + elif T < self.chroma_len: + n_repeat = int(math.ceil(self.chroma_len / T)) + chroma = chroma.repeat(1, n_repeat, 1) + chroma = chroma[:, :self.chroma_len] + logger.debug(f"Chroma was repeated to match length! ({T} -> {chroma.shape[1]})") + + return chroma + + def tokenize(self, x: WavCondition) -> WavCondition: + """Apply WavConditioner tokenization and populate cache if needed.""" + x = super().tokenize(x) + no_undefined_paths = all(p is not None for p in x.path) + if self.cache is not None and no_undefined_paths: + paths = [Path(p) for p in x.path if p is not None] + self.cache.populate_embed_cache(paths, x) + return x + + +class JointEmbeddingConditioner(BaseConditioner): + """Joint embedding conditioning supporting both audio or text conditioning. + + Args: + dim (int): Dimension. + output_dim (int): Output dimension. + device (str): Device. + attribute (str): Attribute used by the conditioner. + autocast_dtype (str): Autocast for the conditioner. + quantize (bool): Whether to quantize the CLAP embedding. + n_q (int): Number of residual quantizers (used if quantize is true). + bins (int): Quantizers' codebooks size (used if quantize is true). + kwargs: Additional parameters for residual vector quantizer. + """ + def __init__(self, dim: int, output_dim: int, device: str, attribute: str, + autocast_dtype: tp.Optional[str] = 'float32', quantize: bool = True, + n_q: int = 12, bins: int = 1024, **kwargs): + super().__init__(dim=dim, output_dim=output_dim) + self.device = device + self.attribute = attribute + if autocast_dtype is None or device == 'cpu': + self.autocast = TorchAutocast(enabled=False) + logger.warning("JointEmbeddingConditioner has no autocast, this might lead to NaN.") + else: + dtype = getattr(torch, autocast_dtype) + assert isinstance(dtype, torch.dtype) + logger.info(f"JointEmbeddingConditioner will be evaluated with autocast as {autocast_dtype}.") + self.autocast = TorchAutocast(enabled=True, device_type=self.device, dtype=dtype) + # residual vector quantizer to discretize the conditioned embedding + self.quantizer: tp.Optional[ResidualVectorQuantizer] = None + if quantize: + self.quantizer = ResidualVectorQuantizer(dim, n_q=n_q, bins=bins, **kwargs) + + def _get_embed(self, x: JointEmbedCondition) -> tp.Tuple[torch.Tensor, torch.Tensor]: + """Get joint embedding in latent space from the inputs. + + Returns: + tuple[torch.Tensor, torch.Tensor]: Tensor for the latent embedding + and corresponding empty indexes. + """ + raise NotImplementedError() + + def forward(self, x: JointEmbedCondition) -> ConditionType: + with self.autocast: + embed, empty_idx = self._get_embed(x) + if self.quantizer is not None: + embed = embed.view(-1, self.dim, 1) + q_res = self.quantizer(embed, frame_rate=1) + out_embed = q_res.x.view(-1, self.dim) + else: + out_embed = embed + out_embed = self.output_proj(out_embed).view(-1, 1, self.output_dim) + mask = torch.ones(*out_embed.shape[:2], device=out_embed.device) + mask[empty_idx, :] = 0 # zero-out index where the input is non-existant + out_embed = (out_embed * mask.unsqueeze(-1)) + return out_embed, mask + + def tokenize(self, x: JointEmbedCondition) -> JointEmbedCondition: + return x + + +class CLAPEmbeddingConditioner(JointEmbeddingConditioner): + """Joint Embedding conditioner based on pre-trained CLAP model. + + This CLAP-based conditioner supports a caching mechanism + over the computed embeddings for faster training. + + Args: + dim (int): Dimension. + output_dim (int): Output dimension. + device (str): Device. + attribute (str): Attribute used by the conditioner. + quantize (bool): Whether to quantize the CLAP embedding. + n_q (int): Number of residual quantizers (used if quantize is true). + bins (int): Quantizers' codebooks size (used if quantize is true). + checkpoint (str): Path to CLAP checkpoint. + model_arch (str): CLAP model architecture. + enable_fusion (bool): Enable fusion for CLAP model. + sample_rate (int): Sample rate used by CLAP model. + max_audio_length (float): Maximum audio length for CLAP model. + audio_stride (float): Stride to use for getting a CLAP embedding on the full sequence. + normalize (bool): Whether to normalize the CLAP embedding. + text_p (float): Probability of using text representation instead of audio at train time. + batch_size (Optional[int]): Batch size for CLAP embedding computation. + autocast_dtype (str): Autocast for the conditioner. + cache_path (Optional[str]): Path for pre-computed embeddings caching. + kwargs: Additional parameters for residual vector quantizer. + """ + def __init__(self, dim: int, output_dim: int, device: str, attribute: str, + quantize: bool, n_q: int, bins: int, checkpoint: tp.Union[str, Path], model_arch: str, + enable_fusion: bool, sample_rate: int, max_audio_length: int, audio_stride: int, + normalize: bool, text_p: bool, batch_size: tp.Optional[int] = None, + autocast_dtype: tp.Optional[str] = 'float32', cache_path: tp.Optional[str] = None, **kwargs): + try: + import laion_clap # type: ignore + except ImportError: + raise ImportError("Please install CLAP to use the CLAPEmbeddingConditioner: 'pip install laion_clap'") + warnings.warn("Sample rate for CLAP conditioner was fixed in version v1.1.0, (from 44.1 to 48 kHz). " + "Please retrain all models.") + checkpoint = AudioCraftEnvironment.resolve_reference_path(checkpoint) + clap_tokenize = RobertaTokenizer.from_pretrained('roberta-base') + clap_model = laion_clap.CLAP_Module(enable_fusion=enable_fusion, amodel=model_arch) + load_clap_state_dict(clap_model, checkpoint) + clap_model.eval() + clap_model.to(device) + super().__init__(dim=dim, output_dim=output_dim, device=device, attribute=attribute, + autocast_dtype=autocast_dtype, quantize=quantize, n_q=n_q, bins=bins, + **kwargs) + self.checkpoint = checkpoint + self.enable_fusion = enable_fusion + self.model_arch = model_arch + self.clap: laion_clap.CLAP_Module + self.clap_tokenize: RobertaTokenizer + self.clap_sample_rate = sample_rate + self.clap_max_frames = int(self.clap_sample_rate * max_audio_length) + self.clap_stride = int(self.clap_sample_rate * audio_stride) + self.batch_size = batch_size or 1 + self.normalize = normalize + self.text_p = text_p + self.__dict__['clap_tokenize'] = clap_tokenize + self.__dict__['clap'] = clap_model + self.wav_cache, self.text_cache = None, None + if cache_path is not None: + self.wav_cache = EmbeddingCache(Path(cache_path) / 'wav', self.device, + compute_embed_fn=self._get_wav_embedding_for_cache, + extract_embed_fn=self._extract_wav_embedding_chunk) + self.text_cache = EmbeddingCache(Path(cache_path) / 'text', self.device, + compute_embed_fn=self._get_text_embedding_for_cache) + + def _tokenizer(self, texts: tp.Union[str, tp.List[str]]) -> dict: + # we use the default params from CLAP module here as well + return self.clap_tokenize(texts, padding="max_length", truncation=True, max_length=77, return_tensors="pt") + + def _compute_text_embedding(self, text: tp.List[str]) -> torch.Tensor: + """Compute text embedding from CLAP model on a given a batch of text. + + Args: + text (list[str]): List of text for the batch, with B items. + Returns: + torch.Tensor: CLAP embedding derived from text, of shape [B, 1, D], with D the CLAP embedding dimension. + """ + with torch.no_grad(): + embed = self.clap.get_text_embedding(text, tokenizer=self._tokenizer, use_tensor=True) + return embed.view(embed.size(0), 1, embed.size(-1)) + + def _get_text_embedding_for_cache(self, path: tp.Union[Path, str], + x: JointEmbedCondition, idx: int) -> torch.Tensor: + """Get text embedding function for the cache.""" + text = x.text[idx] + text = text if text is not None else "" + return self._compute_text_embedding([text])[0] + + def _preprocess_wav(self, wav: torch.Tensor, length: torch.Tensor, sample_rates: tp.List[int]) -> torch.Tensor: + """Preprocess wav to expected format by CLAP model. + + Args: + wav (torch.Tensor): Audio wav, of shape [B, C, T]. + length (torch.Tensor): Actual length of the audio for each item in the batch, of shape [B]. + sample_rates (list[int]): Sample rates for each sample in the batch + Returns: + torch.Tensor: Audio wav of shape [B, T]. + """ + assert wav.dim() == 3, "Expecting wav to be [B, C, T]" + if sample_rates is not None: + _wav = [] + for i, audio in enumerate(wav): + sr = sample_rates[i] + audio = convert_audio(audio, from_rate=sr, to_rate=self.clap_sample_rate, to_channels=1) + _wav.append(audio) + wav = torch.stack(_wav, dim=0) + wav = wav.mean(dim=1) + return wav + + def _compute_wav_embedding(self, wav: torch.Tensor, length: torch.Tensor, + sample_rates: tp.List[int], reduce_mean: bool = False) -> torch.Tensor: + """Compute audio wave embedding from CLAP model. + + Since CLAP operates on a fixed sequence length audio inputs and we need to process longer audio sequences, + we calculate the wav embeddings on `clap_max_frames` windows with `clap_stride`-second stride and + average the resulting embeddings. + + Args: + wav (torch.Tensor): Audio wav, of shape [B, C, T]. + length (torch.Tensor): Actual length of the audio for each item in the batch, of shape [B]. + sample_rates (list[int]): Sample rates for each sample in the batch. + reduce_mean (bool): Whether to get the average tensor. + Returns: + torch.Tensor: Audio embedding of shape [B, F, D], F being the number of chunks, D the dimension. + """ + with torch.no_grad(): + wav = self._preprocess_wav(wav, length, sample_rates) + B, T = wav.shape + if T >= self.clap_max_frames: + wav = wav.unfold(-1, self.clap_max_frames, self.clap_stride) # [B, F, T] + else: + wav = wav.view(-1, 1, T) # [B, F, T] with F=1 + wav = einops.rearrange(wav, 'b f t -> (b f) t') + embed_list = [] + for i in range(0, wav.size(0), self.batch_size): + _wav = wav[i:i+self.batch_size, ...] + _embed = self.clap.get_audio_embedding_from_data(_wav, use_tensor=True) + embed_list.append(_embed) + embed = torch.cat(embed_list, dim=0) + embed = einops.rearrange(embed, '(b f) d -> b f d', b=B) + if reduce_mean: + embed = embed.mean(dim=1, keepdim=True) + return embed # [B, F, D] with F=1 if reduce_mean is True + + def _get_wav_embedding_for_cache(self, path: tp.Union[str, Path], + x: JointEmbedCondition, idx: int) -> torch.Tensor: + """Compute audio wave embedding for the cache. + The embedding is computed on a given audio read from file. + + Args: + path (str or Path): Path to the full audio file. + Returns: + torch.Tensor: Single-item tensor of shape [F, D], F being the number of chunks, D the dimension. + """ + wav, sr = audio_read(path) # [C, T] + wav = wav.unsqueeze(0).to(self.device) # [1, C, T] + wav_len = torch.LongTensor([wav.shape[-1]]).to(self.device) + embed = self._compute_wav_embedding(wav, wav_len, [sr], reduce_mean=False) # [B, F, D] + return embed.squeeze(0) # [F, D] + + def _extract_wav_embedding_chunk(self, full_embed: torch.Tensor, x: JointEmbedCondition, idx: int) -> torch.Tensor: + """Extract the chunk of embedding matching the seek_time and length from the full CLAP audio embedding. + + Args: + full_embed (torch.Tensor): CLAP embedding computed on the full wave, of shape [F, D]. + x (JointEmbedCondition): Joint embedding condition for the full batch. + idx (int): Index considered for the given embedding to extract. + Returns: + torch.Tensor: Wav embedding averaged on sliding window, of shape [1, D]. + """ + sample_rate = x.sample_rate[idx] + seek_time = x.seek_time[idx] + seek_time = 0. if seek_time is None else seek_time + clap_stride = int(self.clap_stride / self.clap_sample_rate) * sample_rate + end_seek_time = seek_time + self.clap_max_frames / self.clap_sample_rate + start_offset = int(seek_time * sample_rate // clap_stride) + end_offset = int(end_seek_time * sample_rate // clap_stride) + wav_embed = full_embed[start_offset:end_offset, ...] + wav_embed = wav_embed.mean(dim=0, keepdim=True) + return wav_embed.to(self.device) # [F, D] + + def _get_text_embedding(self, x: JointEmbedCondition) -> torch.Tensor: + """Get CLAP embedding from a batch of text descriptions.""" + no_nullified_cond = x.wav.shape[-1] > 1 # we don't want to read from cache when condition dropout + if self.text_cache is not None and no_nullified_cond: + assert all(p is not None for p in x.path), "Cache requires all JointEmbedCondition paths to be provided" + paths = [Path(p) for p in x.path if p is not None] + embed = self.text_cache.get_embed_from_cache(paths, x) + else: + text = [xi if xi is not None else "" for xi in x.text] + embed = self._compute_text_embedding(text) + if self.normalize: + embed = torch.nn.functional.normalize(embed, p=2.0, dim=-1) + return embed + + def _get_wav_embedding(self, x: JointEmbedCondition) -> torch.Tensor: + """Get CLAP embedding from a batch of audio tensors (and corresponding sample rates).""" + no_undefined_paths = all(p is not None for p in x.path) + no_nullified_cond = x.wav.shape[-1] > 1 # we don't want to read from cache when condition dropout + if self.wav_cache is not None and no_undefined_paths and no_nullified_cond: + paths = [Path(p) for p in x.path if p is not None] + embed = self.wav_cache.get_embed_from_cache(paths, x) + else: + embed = self._compute_wav_embedding(x.wav, x.length, x.sample_rate, reduce_mean=True) + if self.normalize: + embed = torch.nn.functional.normalize(embed, p=2.0, dim=-1) + return embed + + def tokenize(self, x: JointEmbedCondition) -> JointEmbedCondition: + # Trying to limit as much as possible sync points when the cache is warm. + no_undefined_paths = all(p is not None for p in x.path) + if self.wav_cache is not None and no_undefined_paths: + assert all([p is not None for p in x.path]), "Cache requires all JointEmbedCondition paths to be provided" + paths = [Path(p) for p in x.path if p is not None] + self.wav_cache.populate_embed_cache(paths, x) + if self.text_cache is not None and no_undefined_paths: + assert all([p is not None for p in x.path]), "Cache requires all JointEmbedCondition paths to be provided" + paths = [Path(p) for p in x.path if p is not None] + self.text_cache.populate_embed_cache(paths, x) + return x + + def _get_embed(self, x: JointEmbedCondition) -> tp.Tuple[torch.Tensor, torch.Tensor]: + """Extract shared latent representation from either the wav or the text using CLAP.""" + # decide whether to use text embedding at train time or not + use_text_embed = random.random() < self.text_p + if self.training and not use_text_embed: + embed = self._get_wav_embedding(x) + empty_idx = torch.LongTensor([]) # we assume we always have the audio wav + else: + embed = self._get_text_embedding(x) + empty_idx = torch.LongTensor([i for i, xi in enumerate(x.text) if xi is None or xi == ""]) + return embed, empty_idx + + +def dropout_condition(sample: ConditioningAttributes, condition_type: str, condition: str) -> ConditioningAttributes: + """Utility function for nullifying an attribute inside an ConditioningAttributes object. + If the condition is of type "wav", then nullify it using `nullify_condition` function. + If the condition is of any other type, set its value to None. + Works in-place. + """ + if condition_type not in ['text', 'wav', 'joint_embed']: + raise ValueError( + "dropout_condition got an unexpected condition type!" + f" expected 'text', 'wav' or 'joint_embed' but got '{condition_type}'" + ) + + if condition not in getattr(sample, condition_type): + raise ValueError( + "dropout_condition received an unexpected condition!" + f" expected wav={sample.wav.keys()} and text={sample.text.keys()}" + f" but got '{condition}' of type '{condition_type}'!" + ) + + if condition_type == 'wav': + wav_cond = sample.wav[condition] + sample.wav[condition] = nullify_wav(wav_cond) + elif condition_type == 'joint_embed': + embed = sample.joint_embed[condition] + sample.joint_embed[condition] = nullify_joint_embed(embed) + else: + sample.text[condition] = None + + return sample + + +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 AttributeDropout(DropoutModule): + """Dropout with a given probability per attribute. + This is different from the behavior of ClassifierFreeGuidanceDropout as this allows for attributes + to be dropped out separately. For example, "artist" can be dropped while "genre" remains. + This is in contrast to ClassifierFreeGuidanceDropout where if "artist" is dropped "genre" + must also be dropped. + + Args: + p (tp.Dict[str, float]): A dict mapping between attributes and dropout probability. For example: + ... + "genre": 0.1, + "artist": 0.5, + "wav": 0.25, + ... + active_on_eval (bool, optional): Whether the dropout is active at eval. Default to False. + seed (int, optional): Random seed. + """ + def __init__(self, p: tp.Dict[str, tp.Dict[str, float]], active_on_eval: bool = False, seed: int = 1234): + super().__init__(seed=seed) + self.active_on_eval = active_on_eval + # construct dict that return the values from p otherwise 0 + self.p = {} + for condition_type, probs in p.items(): + self.p[condition_type] = defaultdict(lambda: 0, probs) + + def forward(self, samples: tp.List[ConditioningAttributes]) -> tp.List[ConditioningAttributes]: + """ + Args: + samples (list[ConditioningAttributes]): List of conditions. + Returns: + list[ConditioningAttributes]: List of conditions after certain attributes were set to None. + """ + if not self.training and not self.active_on_eval: + return samples + + samples = deepcopy(samples) + for condition_type, ps in self.p.items(): # for condition types [text, wav] + for condition, p in ps.items(): # for attributes of each type (e.g., [artist, genre]) + if torch.rand(1, generator=self.rng).item() < p: + for sample in samples: + dropout_condition(sample, condition_type, condition) + return samples + + def __repr__(self): + return f"AttributeDropout({dict(self.p)})" + + +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: tp.List[ConditioningAttributes]) -> tp.List[ConditioningAttributes]: + """ + Args: + samples (list[ConditioningAttributes]): List of conditions. + Returns: + list[ConditioningAttributes]: List of conditions after all attributes were set to None. + """ + if not self.training: + return samples + + # decide on which attributes to drop in a batched fashion + drop = torch.rand(1, generator=self.rng).item() < self.p + if not drop: + return samples + + # nullify conditions of all attributes + samples = deepcopy(samples) + for condition_type in ["wav", "text"]: + for sample in samples: + for condition in sample.attributes[condition_type]: + dropout_condition(sample, condition_type, condition) + return samples + + def __repr__(self): + return f"ClassifierFreeGuidanceDropout(p={self.p})" + + +class ConditioningProvider(nn.Module): + """Prepare and provide conditions given all the supported conditioners. + + Args: + conditioners (dict): Dictionary of conditioners. + device (torch.device or str, optional): Device for conditioners and output condition types. + """ + def __init__(self, conditioners: tp.Dict[str, BaseConditioner], device: tp.Union[torch.device, str] = "cpu"): + super().__init__() + self.device = device + self.conditioners = nn.ModuleDict(conditioners) + + @property + def joint_embed_conditions(self): + return [m.attribute for m in self.conditioners.values() if isinstance(m, JointEmbeddingConditioner)] + + @property + def has_joint_embed_conditions(self): + return len(self.joint_embed_conditions) > 0 + + @property + def text_conditions(self): + return [k for k, v in self.conditioners.items() if isinstance(v, TextConditioner)] + + @property + def wav_conditions(self): + return [k for k, v in self.conditioners.items() if isinstance(v, WaveformConditioner)] + + @property + def has_wav_condition(self): + return len(self.wav_conditions) > 0 + + def tokenize(self, inputs: tp.List[ConditioningAttributes]) -> tp.Dict[str, tp.Any]: + """Match attributes/wavs with existing conditioners in self, and compute tokenize them accordingly. + This should be called before starting any real GPU work to avoid synchronization points. + This will return a dict matching conditioner names to their arbitrary tokenized representations. + + Args: + inputs (list[ConditioningAttributes]): List of ConditioningAttributes objects containing + text and wav conditions. + """ + assert all([isinstance(x, ConditioningAttributes) for x in inputs]), ( + "Got unexpected types input for conditioner! should be tp.List[ConditioningAttributes]", + f" but types were {set([type(x) for x in inputs])}" + ) + + output = {} + text = self._collate_text(inputs) + wavs = self._collate_wavs(inputs) + joint_embeds = self._collate_joint_embeds(inputs) + + assert set(text.keys() | wavs.keys() | joint_embeds.keys()).issubset(set(self.conditioners.keys())), ( + f"Got an unexpected attribute! Expected {self.conditioners.keys()}, ", + f"got {text.keys(), wavs.keys(), joint_embeds.keys()}" + ) + + for attribute, batch in chain(text.items(), wavs.items(), joint_embeds.items()): + output[attribute] = self.conditioners[attribute].tokenize(batch) + return output + + def forward(self, tokenized: tp.Dict[str, tp.Any]) -> tp.Dict[str, ConditionType]: + """Compute pairs of `(embedding, mask)` using the configured conditioners and the tokenized representations. + The output is for example: + { + "genre": (torch.Tensor([B, 1, D_genre]), torch.Tensor([B, 1])), + "description": (torch.Tensor([B, T_desc, D_desc]), torch.Tensor([B, T_desc])), + ... + } + + Args: + tokenized (dict): Dict of tokenized representations as returned by `tokenize()`. + """ + output = {} + for attribute, inputs in tokenized.items(): + condition, mask = self.conditioners[attribute](inputs) + output[attribute] = (condition, mask) + return output + + def _collate_text(self, samples: tp.List[ConditioningAttributes]) -> tp.Dict[str, tp.List[tp.Optional[str]]]: + """Given a list of ConditioningAttributes objects, compile a dictionary where the keys + are the attributes and the values are the aggregated input per attribute. + For example: + Input: + [ + ConditioningAttributes(text={"genre": "Rock", "description": "A rock song with a guitar solo"}, wav=...), + ConditioningAttributes(text={"genre": "Hip-hop", "description": "A hip-hop verse"}, wav=...), + ] + Output: + { + "genre": ["Rock", "Hip-hop"], + "description": ["A rock song with a guitar solo", "A hip-hop verse"] + } + + Args: + samples (list of ConditioningAttributes): List of ConditioningAttributes samples. + Returns: + dict[str, list[str, optional]]: A dictionary mapping an attribute name to text batch. + """ + out: tp.Dict[str, tp.List[tp.Optional[str]]] = defaultdict(list) + texts = [x.text for x in samples] + for text in texts: + for condition in self.text_conditions: + out[condition].append(text[condition]) + return out + + def _collate_wavs(self, samples: tp.List[ConditioningAttributes]) -> tp.Dict[str, WavCondition]: + """Generate a dict where the keys are attributes by which we fetch similar wavs, + and the values are Tensors of wavs according to said attributes. + + *Note*: by the time the samples reach this function, each sample should have some waveform + inside the "wav" attribute. It should be either: + 1. A real waveform + 2. A null waveform due to the sample having no similar waveforms (nullified by the dataset) + 3. A null waveform due to it being dropped in a dropout module (nullified by dropout) + + Args: + samples (list of ConditioningAttributes): List of ConditioningAttributes samples. + Returns: + dict[str, WavCondition]: A dictionary mapping an attribute name to wavs. + """ + wavs = defaultdict(list) + lengths = defaultdict(list) + sample_rates = defaultdict(list) + paths = defaultdict(list) + seek_times = defaultdict(list) + out: tp.Dict[str, WavCondition] = {} + + for sample in samples: + for attribute in self.wav_conditions: + wav, length, sample_rate, path, seek_time = sample.wav[attribute] + assert wav.dim() == 3, f"Got wav with dim={wav.dim()}, but expected 3 [1, C, T]" + assert wav.size(0) == 1, f"Got wav [B, C, T] with shape={wav.shape}, but expected B == 1" + # mono-channel conditioning + wav = wav.mean(1, keepdim=True) # [1, 1, T] + wavs[attribute].append(wav.flatten()) # [T] + lengths[attribute].append(length) + sample_rates[attribute].extend(sample_rate) + paths[attribute].extend(path) + seek_times[attribute].extend(seek_time) + + # stack all wavs to a single tensor + for attribute in self.wav_conditions: + stacked_wav, _ = collate(wavs[attribute], dim=0) + out[attribute] = WavCondition( + stacked_wav.unsqueeze(1), torch.cat(lengths[attribute]), sample_rates[attribute], + paths[attribute], seek_times[attribute]) + + return out + + def _collate_joint_embeds(self, samples: tp.List[ConditioningAttributes]) -> tp.Dict[str, JointEmbedCondition]: + """Generate a dict where the keys are attributes by which we compute joint embeddings, + and the values are Tensors of pre-computed embeddings and the corresponding text attributes. + + Args: + samples (list[ConditioningAttributes]): List of ConditioningAttributes samples. + Returns: + A dictionary mapping an attribute name to joint embeddings. + """ + texts = defaultdict(list) + wavs = defaultdict(list) + lengths = defaultdict(list) + sample_rates = defaultdict(list) + paths = defaultdict(list) + seek_times = defaultdict(list) + channels: int = 0 + + out = {} + for sample in samples: + for attribute in self.joint_embed_conditions: + wav, text, length, sample_rate, path, seek_time = sample.joint_embed[attribute] + assert wav.dim() == 3 + if channels == 0: + channels = wav.size(1) + else: + assert channels == wav.size(1), "not all audio has same number of channels in batch" + assert wav.size(0) == 1, "Expecting single-wav batch in the collate method" + wav = einops.rearrange(wav, "b c t -> (b c t)") # [1, C, T] => [C * T] + wavs[attribute].append(wav) + texts[attribute].extend(text) + lengths[attribute].append(length) + sample_rates[attribute].extend(sample_rate) + paths[attribute].extend(path) + seek_times[attribute].extend(seek_time) + + for attribute in self.joint_embed_conditions: + stacked_texts = texts[attribute] + stacked_paths = paths[attribute] + stacked_seek_times = seek_times[attribute] + stacked_wavs = pad_sequence(wavs[attribute]).to(self.device) + stacked_wavs = einops.rearrange(stacked_wavs, "(c t) b -> b c t", c=channels) + stacked_sample_rates = sample_rates[attribute] + stacked_lengths = torch.cat(lengths[attribute]).to(self.device) + assert stacked_lengths.size(0) == stacked_wavs.size(0) + assert len(stacked_sample_rates) == stacked_wavs.size(0) + assert len(stacked_texts) == stacked_wavs.size(0) + out[attribute] = JointEmbedCondition( + text=stacked_texts, wav=stacked_wavs, + length=stacked_lengths, sample_rate=stacked_sample_rates, + path=stacked_paths, seek_time=stacked_seek_times) + + return out + + +class ConditionFuser(StreamingModule): + """Condition fuser handles the logic to combine the different conditions + to the actual model input. + + Args: + fuse2cond (tp.Dict[str, str]): A dictionary that says how to fuse + each condition. For example: + { + "prepend": ["description"], + "sum": ["genre", "bpm"], + "cross": ["description"], + } + cross_attention_pos_emb (bool, optional): Use positional embeddings in cross attention. + cross_attention_pos_emb_scale (int): Scale for positional embeddings in cross attention if used. + """ + FUSING_METHODS = ["sum", "prepend", "cross", "input_interpolate"] + + def __init__(self, fuse2cond: tp.Dict[str, tp.List[str]], cross_attention_pos_emb: bool = False, + cross_attention_pos_emb_scale: float = 1.0): + super().__init__() + assert all( + [k in self.FUSING_METHODS for k in fuse2cond.keys()] + ), f"Got invalid fuse method, allowed methods: {self.FUSING_METHODS}" + self.cross_attention_pos_emb = cross_attention_pos_emb + self.cross_attention_pos_emb_scale = cross_attention_pos_emb_scale + self.fuse2cond: tp.Dict[str, tp.List[str]] = fuse2cond + self.cond2fuse: tp.Dict[str, str] = {} + for fuse_method, conditions in fuse2cond.items(): + for condition in conditions: + self.cond2fuse[condition] = fuse_method + + def forward( + self, + input: torch.Tensor, + conditions: tp.Dict[str, ConditionType] + ) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]: + """Fuse the conditions to the provided model input. + + Args: + input (torch.Tensor): Transformer input. + conditions (dict[str, ConditionType]): Dict of conditions. + Returns: + tuple[torch.Tensor, torch.Tensor]: The first tensor is the transformer input + after the conditions have been fused. The second output tensor is the tensor + used for cross-attention or None if no cross attention inputs exist. + """ + B, T, _ = input.shape + + if 'offsets' in self._streaming_state: + first_step = False + offsets = self._streaming_state['offsets'] + else: + first_step = True + offsets = torch.zeros(input.shape[0], dtype=torch.long, device=input.device) + + assert set(conditions.keys()).issubset(set(self.cond2fuse.keys())), \ + f"given conditions contain unknown attributes for fuser, " \ + f"expected {self.cond2fuse.keys()}, got {conditions.keys()}" + cross_attention_output = None + for cond_type, (cond, cond_mask) in conditions.items(): + op = self.cond2fuse[cond_type] + if op == 'sum': + input += cond + elif op == 'input_interpolate': + cond = einops.rearrange(cond, "b t d -> b d t") + cond = F.interpolate(cond, size=input.shape[1]) + input += einops.rearrange(cond, "b d t -> b t d") + elif op == 'prepend': + if first_step: + input = torch.cat([cond, input], dim=1) + elif op == 'cross': + if cross_attention_output is not None: + cross_attention_output = torch.cat([cross_attention_output, cond], dim=1) + else: + cross_attention_output = cond + else: + raise ValueError(f"unknown op ({op})") + + if self.cross_attention_pos_emb and cross_attention_output is not None: + positions = torch.arange( + cross_attention_output.shape[1], + device=cross_attention_output.device + ).view(1, -1, 1) + pos_emb = create_sin_embedding(positions, cross_attention_output.shape[-1]) + cross_attention_output = cross_attention_output + self.cross_attention_pos_emb_scale * pos_emb + + if self._is_streaming: + self._streaming_state['offsets'] = offsets + T + + return input, cross_attention_output diff --git a/src/audiocraft/modules/conv.py b/src/audiocraft/modules/conv.py new file mode 100644 index 0000000000000000000000000000000000000000..d115cbf8729b642ed78608bd00a4d0fd5afae6fd --- /dev/null +++ b/src/audiocraft/modules/conv.py @@ -0,0 +1,243 @@ +# 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 warnings + +import torch +from torch import nn +from torch.nn import functional as F +from torch.nn.utils import spectral_norm, weight_norm + + +CONV_NORMALIZATIONS = frozenset(['none', 'weight_norm', 'spectral_norm', + 'time_group_norm']) + + +def apply_parametrization_norm(module: nn.Module, norm: str = 'none'): + 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): + """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 == '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 = 'constant', 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): + x = self.conv(x) + x = self.norm(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): + x = self.conv(x) + x = self.norm(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): + x = self.convtr(x) + x = self.norm(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): + x = self.convtr(x) + x = self.norm(x) + return x + + +class StreamableConv1d(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("StreamableConv1d 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): + 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] + kernel_size = (kernel_size - 1) * dilation + 1 # effective kernel size with dilations + padding_total = kernel_size - stride + 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) + return self.conv(x) + + +class StreamableConvTranspose1d(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): + 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)) + return y diff --git a/src/audiocraft/modules/diffusion_schedule.py b/src/audiocraft/modules/diffusion_schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..74ca6e3f2e7c4ff904d96dade315b0b46856778d --- /dev/null +++ b/src/audiocraft/modules/diffusion_schedule.py @@ -0,0 +1,272 @@ +# 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. + +""" +Functions for Noise Schedule, defines diffusion process, reverse process and data processor. +""" + +from collections import namedtuple +import random +import typing as tp +import julius +import torch + +TrainingItem = namedtuple("TrainingItem", "noisy noise step") + + +def betas_from_alpha_bar(alpha_bar): + alphas = torch.cat([torch.Tensor([alpha_bar[0]]), alpha_bar[1:]/alpha_bar[:-1]]) + return 1 - alphas + + +class SampleProcessor(torch.nn.Module): + def project_sample(self, x: torch.Tensor): + """Project the original sample to the 'space' where the diffusion will happen.""" + return x + + def return_sample(self, z: torch.Tensor): + """Project back from diffusion space to the actual sample space.""" + return z + + +class MultiBandProcessor(SampleProcessor): + """ + MultiBand sample processor. The input audio is splitted across + frequency bands evenly distributed in mel-scale. + + Each band will be rescaled to match the power distribution + of Gaussian noise in that band, using online metrics + computed on the first few samples. + + Args: + n_bands (int): Number of mel-bands to split the signal over. + sample_rate (int): Sample rate of the audio. + num_samples (int): Number of samples to use to fit the rescaling + for each band. The processor won't be stable + until it has seen that many samples. + power_std (float or list/tensor): The rescaling factor computed to match the + power of Gaussian noise in each band is taken to + that power, i.e. `1.` means full correction of the energy + in each band, and values less than `1` means only partial + correction. Can be used to balance the relative importance + of low vs. high freq in typical audio signals. + """ + def __init__(self, n_bands: int = 8, sample_rate: float = 24_000, + num_samples: int = 10_000, power_std: tp.Union[float, tp.List[float], torch.Tensor] = 1.): + super().__init__() + self.n_bands = n_bands + self.split_bands = julius.SplitBands(sample_rate, n_bands=n_bands) + self.num_samples = num_samples + self.power_std = power_std + if isinstance(power_std, list): + assert len(power_std) == n_bands + power_std = torch.tensor(power_std) + self.register_buffer('counts', torch.zeros(1)) + self.register_buffer('sum_x', torch.zeros(n_bands)) + self.register_buffer('sum_x2', torch.zeros(n_bands)) + self.register_buffer('sum_target_x2', torch.zeros(n_bands)) + self.counts: torch.Tensor + self.sum_x: torch.Tensor + self.sum_x2: torch.Tensor + self.sum_target_x2: torch.Tensor + + @property + def mean(self): + mean = self.sum_x / self.counts + return mean + + @property + def std(self): + std = (self.sum_x2 / self.counts - self.mean**2).clamp(min=0).sqrt() + return std + + @property + def target_std(self): + target_std = self.sum_target_x2 / self.counts + return target_std + + def project_sample(self, x: torch.Tensor): + assert x.dim() == 3 + bands = self.split_bands(x) + if self.counts.item() < self.num_samples: + ref_bands = self.split_bands(torch.randn_like(x)) + self.counts += len(x) + self.sum_x += bands.mean(dim=(2, 3)).sum(dim=1) + self.sum_x2 += bands.pow(2).mean(dim=(2, 3)).sum(dim=1) + self.sum_target_x2 += ref_bands.pow(2).mean(dim=(2, 3)).sum(dim=1) + rescale = (self.target_std / self.std.clamp(min=1e-12)) ** self.power_std # same output size + bands = (bands - self.mean.view(-1, 1, 1, 1)) * rescale.view(-1, 1, 1, 1) + return bands.sum(dim=0) + + def return_sample(self, x: torch.Tensor): + assert x.dim() == 3 + bands = self.split_bands(x) + rescale = (self.std / self.target_std) ** self.power_std + bands = bands * rescale.view(-1, 1, 1, 1) + self.mean.view(-1, 1, 1, 1) + return bands.sum(dim=0) + + +class NoiseSchedule: + """Noise schedule for diffusion. + + Args: + beta_t0 (float): Variance of the first diffusion step. + beta_t1 (float): Variance of the last diffusion step. + beta_exp (float): Power schedule exponent + num_steps (int): Number of diffusion step. + variance (str): choice of the sigma value for the denoising eq. Choices: "beta" or "beta_tilde" + clip (float): clipping value for the denoising steps + rescale (float): rescaling value to avoid vanishing signals unused by default (i.e 1) + repartition (str): shape of the schedule only power schedule is supported + sample_processor (SampleProcessor): Module that normalize data to match better the gaussian distribution + noise_scale (float): Scaling factor for the noise + """ + def __init__(self, beta_t0: float = 1e-4, beta_t1: float = 0.02, num_steps: int = 1000, variance: str = 'beta', + clip: float = 5., rescale: float = 1., device='cuda', beta_exp: float = 1, + repartition: str = "power", alpha_sigmoid: dict = {}, n_bands: tp.Optional[int] = None, + sample_processor: SampleProcessor = SampleProcessor(), noise_scale: float = 1.0, **kwargs): + + self.beta_t0 = beta_t0 + self.beta_t1 = beta_t1 + self.variance = variance + self.num_steps = num_steps + self.clip = clip + self.sample_processor = sample_processor + self.rescale = rescale + self.n_bands = n_bands + self.noise_scale = noise_scale + assert n_bands is None + if repartition == "power": + self.betas = torch.linspace(beta_t0 ** (1 / beta_exp), beta_t1 ** (1 / beta_exp), num_steps, + device=device, dtype=torch.float) ** beta_exp + else: + raise RuntimeError('Not implemented') + self.rng = random.Random(1234) + + def get_beta(self, step: tp.Union[int, torch.Tensor]): + if self.n_bands is None: + return self.betas[step] + else: + return self.betas[:, step] # [n_bands, len(step)] + + def get_initial_noise(self, x: torch.Tensor): + if self.n_bands is None: + return torch.randn_like(x) + return torch.randn((x.size(0), self.n_bands, x.size(2))) + + def get_alpha_bar(self, step: tp.Optional[tp.Union[int, torch.Tensor]] = None) -> torch.Tensor: + """Return 'alpha_bar', either for a given step, or as a tensor with its value for each step.""" + if step is None: + return (1 - self.betas).cumprod(dim=-1) # works for simgle and multi bands + if type(step) is int: + return (1 - self.betas[:step + 1]).prod() + else: + return (1 - self.betas).cumprod(dim=0)[step].view(-1, 1, 1) + + def get_training_item(self, x: torch.Tensor, tensor_step: bool = False) -> TrainingItem: + """Create a noisy data item for diffusion model training: + + Args: + x (torch.Tensor): clean audio data torch.tensor(bs, 1, T) + tensor_step (bool): If tensor_step = false, only one step t is sample, + the whole batch is diffused to the same step and t is int. + If tensor_step = true, t is a tensor of size (x.size(0),) + every element of the batch is diffused to a independently sampled. + """ + step: tp.Union[int, torch.Tensor] + if tensor_step: + bs = x.size(0) + step = torch.randint(0, self.num_steps, size=(bs,), device=x.device) + else: + step = self.rng.randrange(self.num_steps) + alpha_bar = self.get_alpha_bar(step) # [batch_size, n_bands, 1] + + x = self.sample_processor.project_sample(x) + noise = torch.randn_like(x) + noisy = (alpha_bar.sqrt() / self.rescale) * x + (1 - alpha_bar).sqrt() * noise * self.noise_scale + return TrainingItem(noisy, noise, step) + + def generate(self, model: torch.nn.Module, initial: tp.Optional[torch.Tensor] = None, + condition: tp.Optional[torch.Tensor] = None, return_list: bool = False): + """Full ddpm reverse process. + + Args: + model (nn.Module): Diffusion model. + initial (tensor): Initial Noise. + condition (tensor): Input conditionning Tensor (e.g. encodec compressed representation). + return_list (bool): Whether to return the whole process or only the sampled point. + """ + alpha_bar = self.get_alpha_bar(step=self.num_steps - 1) + current = initial + iterates = [initial] + for step in range(self.num_steps)[::-1]: + with torch.no_grad(): + estimate = model(current, step, condition=condition).sample + alpha = 1 - self.betas[step] + previous = (current - (1 - alpha) / (1 - alpha_bar).sqrt() * estimate) / alpha.sqrt() + previous_alpha_bar = self.get_alpha_bar(step=step - 1) + if step == 0: + sigma2 = 0 + elif self.variance == 'beta': + sigma2 = 1 - alpha + elif self.variance == 'beta_tilde': + sigma2 = (1 - previous_alpha_bar) / (1 - alpha_bar) * (1 - alpha) + elif self.variance == 'none': + sigma2 = 0 + else: + raise ValueError(f'Invalid variance type {self.variance}') + + if sigma2 > 0: + previous += sigma2**0.5 * torch.randn_like(previous) * self.noise_scale + if self.clip: + previous = previous.clamp(-self.clip, self.clip) + current = previous + alpha_bar = previous_alpha_bar + if step == 0: + previous *= self.rescale + if return_list: + iterates.append(previous.cpu()) + + if return_list: + return iterates + else: + return self.sample_processor.return_sample(previous) + + def generate_subsampled(self, model: torch.nn.Module, initial: torch.Tensor, step_list: tp.Optional[list] = None, + condition: tp.Optional[torch.Tensor] = None, return_list: bool = False): + """Reverse process that only goes through Markov chain states in step_list.""" + if step_list is None: + step_list = list(range(1000))[::-50] + [0] + alpha_bar = self.get_alpha_bar(step=self.num_steps - 1) + alpha_bars_subsampled = (1 - self.betas).cumprod(dim=0)[list(reversed(step_list))].cpu() + betas_subsampled = betas_from_alpha_bar(alpha_bars_subsampled) + current = initial * self.noise_scale + iterates = [current] + for idx, step in enumerate(step_list[:-1]): + with torch.no_grad(): + estimate = model(current, step, condition=condition).sample * self.noise_scale + alpha = 1 - betas_subsampled[-1 - idx] + previous = (current - (1 - alpha) / (1 - alpha_bar).sqrt() * estimate) / alpha.sqrt() + previous_alpha_bar = self.get_alpha_bar(step_list[idx + 1]) + if step == step_list[-2]: + sigma2 = 0 + previous_alpha_bar = torch.tensor(1.0) + else: + sigma2 = (1 - previous_alpha_bar) / (1 - alpha_bar) * (1 - alpha) + if sigma2 > 0: + previous += sigma2**0.5 * torch.randn_like(previous) * self.noise_scale + if self.clip: + previous = previous.clamp(-self.clip, self.clip) + current = previous + alpha_bar = previous_alpha_bar + if step == 0: + previous *= self.rescale + if return_list: + iterates.append(previous.cpu()) + if return_list: + return iterates + else: + return self.sample_processor.return_sample(previous) diff --git a/src/audiocraft/modules/lstm.py b/src/audiocraft/modules/lstm.py new file mode 100644 index 0000000000000000000000000000000000000000..c0866175950c1ca4f6cca98649525e6481853bba --- /dev/null +++ b/src/audiocraft/modules/lstm.py @@ -0,0 +1,25 @@ +# 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 torch import nn + + +class StreamableLSTM(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 diff --git a/src/audiocraft/modules/rope.py b/src/audiocraft/modules/rope.py new file mode 100644 index 0000000000000000000000000000000000000000..c12cee0954f27c45d79627771fdf7fa9fc10dfcc --- /dev/null +++ b/src/audiocraft/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 diff --git a/src/audiocraft/modules/seanet.py b/src/audiocraft/modules/seanet.py new file mode 100644 index 0000000000000000000000000000000000000000..3e5998e9153afb6e68ea410d565e00ea835db248 --- /dev/null +++ b/src/audiocraft/modules/seanet.py @@ -0,0 +1,258 @@ +# 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 + +import numpy as np +import torch.nn as nn + +from .conv import StreamableConv1d, StreamableConvTranspose1d +from .lstm import StreamableLSTM + + +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 + (streamable) 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 = 'none', 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)): + in_chs = dim if i == 0 else hidden + out_chs = dim if i == len(kernel_sizes) - 1 else hidden + block += [ + act(**activation_params), + StreamableConv1d(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 + if true_skip: + self.shortcut = nn.Identity() + else: + self.shortcut = StreamableConv1d(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) + + +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. We use the decoder order as some models may only employ the decoder. + 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. + 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. + disable_norm_outer_blocks (int): Number of blocks for which we don't apply norm. + For the encoder, it corresponds to the N first blocks. + """ + def __init__(self, channels: int = 1, dimension: int = 128, n_filters: int = 32, n_residual_layers: int = 3, + ratios: tp.List[int] = [8, 5, 4, 2], activation: str = 'ELU', activation_params: dict = {'alpha': 1.0}, + norm: str = 'none', 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 = True, compress: int = 2, lstm: int = 0, + disable_norm_outer_blocks: int = 0): + 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) + self.n_blocks = len(self.ratios) + 2 # first and last conv + residual blocks + self.disable_norm_outer_blocks = disable_norm_outer_blocks + assert self.disable_norm_outer_blocks >= 0 and self.disable_norm_outer_blocks <= self.n_blocks, \ + "Number of blocks for which to disable norm is invalid." \ + "It should be lower or equal to the actual number of blocks in the network and greater or equal to 0." + + act = getattr(nn, activation) + mult = 1 + model: tp.List[nn.Module] = [ + StreamableConv1d(channels, mult * n_filters, kernel_size, + norm='none' if self.disable_norm_outer_blocks >= 1 else norm, + norm_kwargs=norm_params, causal=causal, pad_mode=pad_mode) + ] + # Downsample to raw audio scale + for i, ratio in enumerate(self.ratios): + block_norm = 'none' if self.disable_norm_outer_blocks >= i + 2 else norm + # Add residual layers + for j in range(n_residual_layers): + model += [ + SEANetResnetBlock(mult * n_filters, kernel_sizes=[residual_kernel_size, 1], + dilations=[dilation_base ** j, 1], + norm=block_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), + StreamableConv1d(mult * n_filters, mult * n_filters * 2, + kernel_size=ratio * 2, stride=ratio, + norm=block_norm, norm_kwargs=norm_params, + causal=causal, pad_mode=pad_mode), + ] + mult *= 2 + + if lstm: + model += [StreamableLSTM(mult * n_filters, num_layers=lstm)] + + model += [ + act(**activation_params), + StreamableConv1d(mult * n_filters, dimension, last_kernel_size, + norm='none' if self.disable_norm_outer_blocks == self.n_blocks else 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. + disable_norm_outer_blocks (int): Number of blocks for which we don't apply norm. + For the decoder, it corresponds to the N last blocks. + 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 = 3, + 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 = 'none', 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 = True, compress: int = 2, lstm: int = 0, + disable_norm_outer_blocks: int = 0, 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) + self.n_blocks = len(self.ratios) + 2 # first and last conv + residual blocks + self.disable_norm_outer_blocks = disable_norm_outer_blocks + assert self.disable_norm_outer_blocks >= 0 and self.disable_norm_outer_blocks <= self.n_blocks, \ + "Number of blocks for which to disable norm is invalid." \ + "It should be lower or equal to the actual number of blocks in the network and greater or equal to 0." + + act = getattr(nn, activation) + mult = int(2 ** len(self.ratios)) + model: tp.List[nn.Module] = [ + StreamableConv1d(dimension, mult * n_filters, kernel_size, + norm='none' if self.disable_norm_outer_blocks == self.n_blocks else norm, + norm_kwargs=norm_params, causal=causal, pad_mode=pad_mode) + ] + + if lstm: + model += [StreamableLSTM(mult * n_filters, num_layers=lstm)] + + # Upsample to raw audio scale + for i, ratio in enumerate(self.ratios): + block_norm = 'none' if self.disable_norm_outer_blocks >= self.n_blocks - (i + 1) else norm + # Add upsampling layers + model += [ + act(**activation_params), + StreamableConvTranspose1d(mult * n_filters, mult * n_filters // 2, + kernel_size=ratio * 2, stride=ratio, + norm=block_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=block_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), + StreamableConv1d(n_filters, channels, last_kernel_size, + norm='none' if self.disable_norm_outer_blocks >= 1 else 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: + 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 diff --git a/src/audiocraft/modules/streaming.py b/src/audiocraft/modules/streaming.py new file mode 100644 index 0000000000000000000000000000000000000000..fba06936294ca15d72acd2d44f9dbda39a638107 --- /dev/null +++ b/src/audiocraft/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 diff --git a/src/audiocraft/modules/transformer.py b/src/audiocraft/modules/transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..4d44b39ecbe6ce2ec4370a149d0b285ebf663f44 --- /dev/null +++ b/src/audiocraft/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 diff --git a/src/audiocraft/py.typed b/src/audiocraft/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/audiocraft/quantization/__init__.py b/src/audiocraft/quantization/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1e0c7e429ab96d67be667e23bf7a0ffa389c036b --- /dev/null +++ b/src/audiocraft/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 diff --git a/src/audiocraft/quantization/base.py b/src/audiocraft/quantization/base.py new file mode 100644 index 0000000000000000000000000000000000000000..a77fefb98e62a5bbc6385910261ffdde2ffa5a25 --- /dev/null +++ b/src/audiocraft/quantization/base.py @@ -0,0 +1,99 @@ +# 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: + x: torch.Tensor + codes: torch.Tensor + bandwidth: torch.Tensor # bandwidth in kb/s used, per batch item. + penalty: 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") diff --git a/src/audiocraft/quantization/core_vq.py b/src/audiocraft/quantization/core_vq.py new file mode 100644 index 0000000000000000000000000000000000000000..6aaa3b077c53b413e2b2a904ac7e769d1c623b36 --- /dev/null +++ b/src/audiocraft/quantization/core_vq.py @@ -0,0 +1,405 @@ +# 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, 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(), 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 + residual = x + + all_losses = [] + all_indices = [] + + n_q = n_q or len(self.layers) + + for i, layer in enumerate(self.layers[:n_q]): + quantized, indices, loss = layer(residual) + quantized = quantized.detach() + residual = residual - quantized + quantized_out = quantized_out + quantized + all_indices.append(indices) + all_losses.append(loss) + + if self.training: + # Solving subtle bug with STE and RVQ: https://github.com/facebookresearch/encodec/issues/25 + quantized_out = x + (quantized_out - x).detach() + + out_losses, out_indices = map(torch.stack, (all_losses, all_indices)) + return 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 diff --git a/src/audiocraft/quantization/vq.py b/src/audiocraft/quantization/vq.py new file mode 100644 index 0000000000000000000000000000000000000000..aa57bea59db95ddae35e0657f723ca3a29ee943b --- /dev/null +++ b/src/audiocraft/quantization/vq.py @@ -0,0 +1,115 @@ +# 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 + 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, 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 diff --git a/src/audiocraft/utils/__init__.py b/src/audiocraft/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..75e25a0212f98e4a18d97c86c6cda225636a3215 --- /dev/null +++ b/src/audiocraft/utils/__init__.py @@ -0,0 +1,6 @@ +# 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. +"""Utilities.""" diff --git a/src/audiocraft/utils/autocast.py b/src/audiocraft/utils/autocast.py new file mode 100644 index 0000000000000000000000000000000000000000..ed644843bb37cf8a92a20fbd51d6cebaa43b9a08 --- /dev/null +++ b/src/audiocraft/utils/autocast.py @@ -0,0 +1,40 @@ +# 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 torch + + +class TorchAutocast: + """TorchAutocast utility class. + Allows you to enable and disable autocast. This is specially useful + when dealing with different architectures and clusters with different + levels of support. + + Args: + enabled (bool): Whether to enable torch.autocast or not. + args: Additional args for torch.autocast. + kwargs: Additional kwargs for torch.autocast + """ + def __init__(self, enabled: bool, *args, **kwargs): + self.autocast = torch.autocast(*args, **kwargs) if enabled else None + + def __enter__(self): + if self.autocast is None: + return + try: + self.autocast.__enter__() + except RuntimeError: + device = self.autocast.device + dtype = self.autocast.fast_dtype + raise RuntimeError( + f"There was an error autocasting with dtype={dtype} device={device}\n" + "If you are on the FAIR Cluster, you might need to use autocast_dtype=float16" + ) + + def __exit__(self, *args, **kwargs): + if self.autocast is None: + return + self.autocast.__exit__(*args, **kwargs) diff --git a/src/audiocraft/utils/cache.py b/src/audiocraft/utils/cache.py new file mode 100644 index 0000000000000000000000000000000000000000..6ba017a761a29c44d3385e0b483877cb4a8d1ec1 --- /dev/null +++ b/src/audiocraft/utils/cache.py @@ -0,0 +1,324 @@ +# 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 concurrent.futures import ThreadPoolExecutor +from collections import deque +from functools import partial +from hashlib import sha1 +import logging +from pathlib import Path +import sys +import typing as tp +import zipfile + +import flashy +import torch + + +logger = logging.getLogger(__name__) + + +def get_full_embed(full_embed: torch.Tensor, x: tp.Any, idx: int, device: tp.Union[str, torch.device]) -> torch.Tensor: + """Utility function for the EmbeddingCache, returning the full embedding without any chunking. + This method can be used in case there is no need in extracting a chunk of the full embedding + read from the cache. + + Args: + full_embed (torch.Tensor): The full embedding. + x (any): Batch object from which the full embedding is derived. + idx (torch.Tensor): Index of object to consider in the batch object. + Returns: + full_embed (torch.Tensor): The full embedding + """ + return full_embed.to(device) + + +class EmbeddingCache: + """Cache around embeddings computation for faster execution. + The EmbeddingCache is storing pre-computed embeddings on disk and provides a simple API + to retrieve the pre-computed embeddings on full inputs and extract only a given chunk + using a user-provided function. When the cache is warm (all embeddings are pre-computed), + the EmbeddingCache allows for faster training as it removes the need of computing the embeddings. + Additionally, it provides in-memory cache around the loaded embeddings to limit IO footprint + and synchronization points in the forward calls. + + Args: + cache_path (Path): Path to folder where all pre-computed embeddings are saved on disk. + device (str or torch.device): Device on which the embedding is returned. + compute_embed_fn (callable[[Path, any, int], torch.Tensor], optional): Function to compute + the embedding from a given object and path. This user provided function can compute the + embedding from the provided object or using the provided path as entry point. The last parameter + specify the index corresponding to the current embedding in the object that can represent batch metadata. + extract_embed_fn (callable[[torch.Tensor, any, int], torch.Tensor], optional): Function to extract + the desired embedding chunk from the full embedding loaded from the cache. The last parameter + specify the index corresponding to the current embedding in the object that can represent batch metadata. + If not specified, will return the full embedding unmodified. + """ + def __init__(self, cache_path: tp.Union[str, Path], device: tp.Union[str, torch.device], + compute_embed_fn: tp.Callable[[Path, tp.Any, int], torch.Tensor], + extract_embed_fn: tp.Optional[tp.Callable[[torch.Tensor, tp.Any, int], torch.Tensor]] = None): + self.cache_path = Path(cache_path) + self.device = device + self._compute_embed_fn = compute_embed_fn + self._extract_embed_fn: tp.Callable[[torch.Tensor, tp.Any, int], torch.Tensor] + if extract_embed_fn is not None: + self._extract_embed_fn = extract_embed_fn + else: + self._extract_embed_fn = partial(get_full_embed, device=device) + if self.cache_path is not None: + self.cache_path.mkdir(exist_ok=True, parents=True) + logger.info(f"Cache instantiated at: {self.cache_path}") + self.pool = ThreadPoolExecutor(8) + self.pool.__enter__() + self._current_batch_cache: dict = {} + self._memory_cache: dict = {} + + def _get_cache_path(self, path: tp.Union[Path, str]): + """Get cache path for the given file path.""" + sig = sha1(str(path).encode()).hexdigest() + return self.cache_path / sig + + @staticmethod + def _get_full_embed_from_cache(cache: Path): + """Loads full pre-computed embedding from the cache.""" + try: + embed = torch.load(cache, 'cpu') + except Exception as exc: + logger.error("Error loading %s: %r", cache, exc) + embed = None + return embed + + def get_embed_from_cache(self, paths: tp.List[Path], x: tp.Any) -> torch.Tensor: + """Get embedding from cache, computing and storing it to cache if not already cached. + The EmbeddingCache first tries to load the embedding from the in-memory cache + containing the pre-computed chunks populated through `populate_embed_cache`. + If not found, the full embedding is computed and stored on disk to be later accessed + to populate the in-memory cache, and the desired embedding chunk is extracted and returned. + + Args: + paths (list[Path or str]): List of paths from where the embeddings can be loaded. + x (any): Object from which the embedding is extracted. + """ + embeds = [] + for idx, path in enumerate(paths): + cache = self._get_cache_path(path) + if cache in self._current_batch_cache: + embed = self._current_batch_cache[cache] + else: + full_embed = self._compute_embed_fn(path, x, idx) + try: + with flashy.utils.write_and_rename(cache, pid=True) as f: + torch.save(full_embed.cpu(), f) + except Exception as exc: + logger.error('Error saving embed %s (%s): %r', cache, full_embed.shape, exc) + else: + logger.info('New embed cache saved: %s (%s)', cache, full_embed.shape) + embed = self._extract_embed_fn(full_embed, x, idx) + embeds.append(embed) + embed = torch.stack(embeds, dim=0) + return embed + + def populate_embed_cache(self, paths: tp.List[Path], x: tp.Any) -> None: + """Populate in-memory caches for embeddings reading from the embeddings stored on disk. + The in-memory caches consist in a cache for the full embedding and another cache for the + final embedding chunk. Such caches are used to limit the IO access when computing the actual embeddings + and reduce the IO footprint and synchronization points during forward passes. + + Args: + paths (list[Path]): List of paths from where the embeddings can be loaded. + x (any): Object from which the embedding is extracted. + """ + self._current_batch_cache.clear() + if self.cache_path is not None: + futures: list = [] + for path in paths: + assert path is not None, "Path is required for computation from cache" + cache = self._get_cache_path(path) + if cache in self._memory_cache or not cache.exists(): + futures.append(None) + else: + futures.append(self.pool.submit(EmbeddingCache._get_full_embed_from_cache, cache)) + for idx, (path, future) in enumerate(zip(paths, futures)): + assert path is not None + cache = self._get_cache_path(path) + full_embed = None + if future is None: + if cache in self._memory_cache: + full_embed = self._memory_cache[cache] + else: + full_embed = future.result() + if full_embed is not None: + self._memory_cache[cache] = full_embed + full_embed = full_embed.to(self.device) + if full_embed is not None: + embed = self._extract_embed_fn(full_embed, x, idx) + self._current_batch_cache[cache] = embed + + +class CachedBatchWriter: + """Write pre computed caches for mini batches. This can + make loading a lot more efficient depending on your filesystem. + + Args: + cache_folder (Path): folder in which the cached minibatches + will be stored. + + Inside cache folder, the structure is the following: + `epoch_number / update_number.zip` + And the zip file contains one entry per batch item. + + It is possible to use the cache with a batch size smaller than + created with but obviously not larger. Make sure to call the + `start_epoch(epoch)` method for indicating changes of epochs. + + See the grid `audiocraft/grids/musicgen/musicgen_warmup_cache.py` + for an example of how to warmup the cache. + """ + def __init__(self, cache_folder: Path): + self.cache_folder = cache_folder + self._current_epoch: tp.Optional[int] = None + self._current_index = 0 + + def start_epoch(self, epoch: int): + """Call at the beginning of each epoch. + """ + self._current_epoch = epoch + self._current_index = 0 + self._zip_path.parent.mkdir(exist_ok=True, parents=True) + + @staticmethod + def _get_zip_path(cache_folder: Path, epoch: int, index: int): + return cache_folder / f"{epoch:05d}" / f"{index:06d}.zip" + + @property + def _zip_path(self): + assert self._current_epoch is not None + return CachedBatchWriter._get_zip_path(self.cache_folder, self._current_epoch, self._current_index) + + def save(self, *content): + """Save one mini batch. This function is distributed-aware + and will automatically merge all the items from the different + workers. + """ + all_contents = [] + for rank in range(flashy.distrib.world_size()): + their_content = flashy.distrib.broadcast_object(content, src=rank) + all_contents.append(their_content) + + if flashy.distrib.is_rank_zero(): + idx = 0 + with flashy.utils.write_and_rename(self._zip_path) as tmp: + with zipfile.ZipFile(tmp, 'w') as zf: + for content in all_contents: + for vals in zip(*content): + with zf.open(f'{idx}', 'w') as f: # type: ignore + torch.save(vals, f) + idx += 1 + flashy.distrib.barrier() + self._current_index += 1 + + +class CachedBatchLoader: + """Loader for cached mini-batches dumped with `CachedBatchWriter`. + + Args: + cache_folder (Path): folder in which the cached minibatches are stored. + batch_size (int): batch size (per GPU) expected. + num_workers (int): number of workers to use for loading. + min_length (int): minimum expected length for each epoch. If some + mini-batches are missing, and error is raised. + + This is iterable just like a regular DataLoader. + """ + + def __init__(self, cache_folder: Path, batch_size: int, + num_workers: int = 10, min_length: int = 1): + self.cache_folder = cache_folder + self.batch_size = batch_size + self.num_workers = num_workers + self.min_length = min_length + self._current_epoch: tp.Optional[int] = None + self.sampler = None # for compatibility with the regular DataLoader + + def __len__(self): + path = CachedBatchWriter._get_zip_path(self.cache_folder, self._current_epoch or 0, 0).parent + return len([p for p in path.iterdir() if p.suffix == ".zip"]) + + def start_epoch(self, epoch: int): + """Call at the beginning of each epoch. + """ + self._current_epoch = epoch + + def _zip_path(self, index: int): + assert self._current_epoch is not None + return CachedBatchWriter._get_zip_path(self.cache_folder, self._current_epoch, index) + + def _load_one(self, index: int): + zip_path = self._zip_path(index) + if not zip_path.exists(): + if index < self.min_length: + raise RuntimeError(f"Cache should have at least {self.min_length} batches, but {index} doesn't exist") + + return None + mode = "rb" if sys.version_info >= (3, 9) else "r" + try: + with zipfile.ZipFile(zip_path, 'r') as zf: + rank = flashy.distrib.rank() + world_size = flashy.distrib.world_size() + root = zipfile.Path(zf) + items = list(root.iterdir()) + total_batch_size = self.batch_size * world_size + if len(items) < total_batch_size: + raise RuntimeError( + f"The cache can handle a max batch size of {len(items)}, " + f"but {total_batch_size} is needed.") + start = rank * self.batch_size + items = items[start: start + self.batch_size] + assert len(items) == self.batch_size + entries = [] + entries = [torch.load(item.open(mode), 'cpu') for item in items] # type: ignore + transposed = zip(*entries) + out = [] + for part in transposed: + assert len(part) > 0 + if isinstance(part[0], torch.Tensor): + out.append(torch.stack(part)) + else: + assert isinstance(part, torch.Tensor) + out.append(part) + return out + except Exception: + logger.error("Error when reading zip path %s", zip_path) + raise + + def __iter__(self): + """This will yields tuples, exactly as provided to the + `CachedBatchWriter.save` method. + """ + pool = ThreadPoolExecutor(self.num_workers) + next_index = 0 + queue = deque() + + def _get_next(): + nonlocal next_index + r = queue.popleft().result() + if r is None: + return None + else: + queue.append(pool.submit(self._load_one, next_index)) + next_index += 1 + return r + + with pool: + # fill the buffer of fetching jobs. + for _ in range(2 * self.num_workers): + queue.append(pool.submit(self._load_one, next_index)) + next_index += 1 + while True: + batch = _get_next() + if batch is None: + return + yield batch diff --git a/src/audiocraft/utils/cluster.py b/src/audiocraft/utils/cluster.py new file mode 100644 index 0000000000000000000000000000000000000000..3380d031739d473fb859c76b9c25350f47fa77e8 --- /dev/null +++ b/src/audiocraft/utils/cluster.py @@ -0,0 +1,75 @@ +# 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. + +""" +Utility functions for SLURM configuration and cluster settings. +""" + +from enum import Enum +import os +import socket +import typing as tp + +import omegaconf + + +class ClusterType(Enum): + AWS = "aws" + FAIR = "fair" + RSC = "rsc" + LOCAL_DARWIN = "darwin" + DEFAULT = "default" # used for any other cluster. + + +def _guess_cluster_type() -> ClusterType: + uname = os.uname() + fqdn = socket.getfqdn() + if uname.sysname == "Linux" and (uname.release.endswith("-aws") or ".ec2" in fqdn): + return ClusterType.AWS + + if fqdn.endswith(".fair"): + return ClusterType.FAIR + + if fqdn.endswith(".facebook.com"): + return ClusterType.RSC + + if uname.sysname == "Darwin": + return ClusterType.LOCAL_DARWIN + + return ClusterType.DEFAULT + + +def get_cluster_type( + cluster_type: tp.Optional[ClusterType] = None, +) -> tp.Optional[ClusterType]: + if cluster_type is None: + return _guess_cluster_type() + + return cluster_type + + +def get_slurm_parameters( + cfg: omegaconf.DictConfig, cluster_type: tp.Optional[ClusterType] = None +) -> omegaconf.DictConfig: + """Update SLURM parameters in configuration based on cluster type. + If the cluster type is not specify, it infers it automatically. + """ + from ..environment import AudioCraftEnvironment + cluster_type = get_cluster_type(cluster_type) + # apply cluster-specific adjustments + if cluster_type == ClusterType.AWS: + cfg["mem_per_gpu"] = None + cfg["constraint"] = None + cfg["setup"] = [] + elif cluster_type == ClusterType.RSC: + cfg["mem_per_gpu"] = None + cfg["setup"] = [] + cfg["constraint"] = None + cfg["partition"] = "learn" + slurm_exclude = AudioCraftEnvironment.get_slurm_exclude() + if slurm_exclude is not None: + cfg["exclude"] = slurm_exclude + return cfg diff --git a/src/audiocraft/utils/utils.py b/src/audiocraft/utils/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2c5799f8bc4ee07dd8d60d6afe67fbc5a6039215 --- /dev/null +++ b/src/audiocraft/utils/utils.py @@ -0,0 +1,298 @@ +# 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 concurrent.futures import ProcessPoolExecutor +from contextlib import contextmanager +from functools import wraps, lru_cache +import hashlib +import json +import logging +from pathlib import Path +import typing as tp + +import flashy +import flashy.distrib +import omegaconf +import torch +from torch.nn.utils.rnn import pad_sequence + + +logger = logging.getLogger(__name__) + + +def model_hash(model: torch.nn.Module) -> str: + """Return a model hash. This should allow us to track regressions in model init + from the logs of past experiments. + """ + hasher = hashlib.sha1() + for p in model.parameters(): + hasher.update(p.data.cpu().numpy().tobytes()) + return hasher.hexdigest() + + +def dict_from_config(cfg: omegaconf.DictConfig) -> dict: + """Convenience function to map an omegaconf configuration to a dictionary. + + Args: + cfg (omegaconf.DictConfig): Original configuration to map to dict. + Returns: + dict: Config as dictionary object. + """ + dct = omegaconf.OmegaConf.to_container(cfg, resolve=True) + assert isinstance(dct, dict) + return dct + + +def random_subset(dataset, max_samples: int, seed: int = 42) -> torch.utils.data.Subset: + if max_samples >= len(dataset): + return dataset + + generator = torch.Generator().manual_seed(seed) + perm = torch.randperm(len(dataset), generator=generator) + return torch.utils.data.Subset(dataset, perm[:max_samples].tolist()) + + +def get_loader(dataset, num_samples: tp.Optional[int], batch_size: int, + num_workers: int, seed: int, **kwargs) -> torch.utils.data.DataLoader: + """Convenience function to load dataset into a dataloader with optional subset sampling. + + Args: + dataset: Dataset to load. + num_samples (Optional[int]): Number of samples to limit subset size. + batch_size (int): Batch size. + num_workers (int): Number of workers for data loading. + seed (int): Random seed. + """ + if num_samples is not None: + dataset = random_subset(dataset, num_samples, seed) + + dataloader = flashy.distrib.loader( + dataset, + batch_size=batch_size, + num_workers=num_workers, + **kwargs + ) + return dataloader + + +def get_dataset_from_loader(dataloader): + dataset = dataloader.dataset + if isinstance(dataset, torch.utils.data.Subset): + return dataset.dataset + else: + return dataset + + +def multinomial(input: torch.Tensor, num_samples: int, replacement=False, *, generator=None): + """torch.multinomial with arbitrary number of dimensions, and number of candidates on the last dimension. + + Args: + input (torch.Tensor): The input tensor containing probabilities. + num_samples (int): Number of samples to draw. + replacement (bool): Whether to draw with replacement or not. + Keywords args: + generator (torch.Generator): A pseudorandom number generator for sampling. + Returns: + torch.Tensor: Last dimension contains num_samples indices + sampled from the multinomial probability distribution + located in the last dimension of tensor input. + """ + input_ = input.reshape(-1, input.shape[-1]) + output_ = torch.multinomial(input_, num_samples=num_samples, replacement=replacement, generator=generator) + output = output_.reshape(*list(input.shape[:-1]), -1) + return output + + +def sample_top_k(probs: torch.Tensor, k: int) -> torch.Tensor: + """Sample next token from top K values along the last dimension of the input probs tensor. + + Args: + probs (torch.Tensor): Input probabilities with token candidates on the last dimension. + k (int): The k in “top-k”. + Returns: + torch.Tensor: Sampled tokens. + """ + top_k_value, _ = torch.topk(probs, k, dim=-1) + min_value_top_k = top_k_value[..., [-1]] + probs *= (probs >= min_value_top_k).float() + probs.div_(probs.sum(dim=-1, keepdim=True)) + next_token = multinomial(probs, num_samples=1) + return next_token + + +def sample_top_p(probs: torch.Tensor, p: float) -> torch.Tensor: + """Sample next token from top P probabilities along the last dimension of the input probs tensor. + + Args: + probs (torch.Tensor): Input probabilities with token candidates on the last dimension. + p (int): The p in “top-p”. + Returns: + torch.Tensor: Sampled tokens. + """ + probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True) + probs_sum = torch.cumsum(probs_sort, dim=-1) + mask = probs_sum - probs_sort > p + probs_sort *= (~mask).float() + probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True)) + next_token = multinomial(probs_sort, num_samples=1) + next_token = torch.gather(probs_idx, -1, next_token) + return next_token + + +class DummyPoolExecutor: + """Dummy pool executor to use when we actually have only 1 worker. + (e.g. instead of ProcessPoolExecutor). + """ + class DummyResult: + def __init__(self, func, *args, **kwargs): + self.func = func + self.args = args + self.kwargs = kwargs + + def result(self): + return self.func(*self.args, **self.kwargs) + + def __init__(self, workers, mp_context=None): + pass + + def submit(self, func, *args, **kwargs): + return DummyPoolExecutor.DummyResult(func, *args, **kwargs) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, exc_tb): + return + + +def get_pool_executor(num_workers: int, mp_context=None): + return ProcessPoolExecutor(num_workers, mp_context) if num_workers > 1 else DummyPoolExecutor(1) + + +def length_to_mask(lengths: torch.Tensor, max_len: tp.Optional[int] = None) -> torch.Tensor: + """Utility function to convert a tensor of sequence lengths to a mask (useful when working on padded sequences). + For example: [3, 5] => [[1, 1, 1, 0, 0], [1, 1, 1, 1, 1]] + + Args: + lengths (torch.Tensor): tensor with lengths + max_len (int): can set the max length manually. Defaults to None. + Returns: + torch.Tensor: mask with 0s where there is pad tokens else 1s + """ + assert len(lengths.shape) == 1, "Length shape should be 1 dimensional." + final_length = lengths.max().item() if not max_len else max_len + final_length = max(final_length, 1) # if all seqs are of len zero we don't want a zero-size tensor + return torch.arange(final_length, device=lengths.device)[None, :] < lengths[:, None] + + +def hash_trick(word: str, vocab_size: int) -> int: + """Hash trick to pair each word with an index + + Args: + word (str): word we wish to convert to an index + vocab_size (int): size of the vocabulary + Returns: + int: index of the word in the embedding LUT + """ + hash = int(hashlib.sha256(word.encode("utf-8")).hexdigest(), 16) + return hash % vocab_size + + +def with_rank_rng(base_seed: int = 1234): + """Decorator for a function so that the function will use a Random Number Generator + whose state depend on the GPU rank. The original RNG state is restored upon returning. + + Args: + base_seed (int): Random seed. + """ + def _decorator(fun: tp.Callable): + @wraps(fun) + def _decorated(*args, **kwargs): + state = torch.get_rng_state() + seed = base_seed ^ flashy.distrib.rank() + torch.manual_seed(seed) + logger.debug('Rank dependent seed set to %d', seed) + try: + return fun(*args, **kwargs) + finally: + torch.set_rng_state(state) + logger.debug('RNG state restored.') + return _decorated + return _decorator + + +def collate(tensors: tp.List[torch.Tensor], dim: int = 0) -> tp.Tuple[torch.Tensor, torch.Tensor]: + """Get a list of tensors and collate them to a single tensor. according to the following logic: + - `dim` specifies the time dimension which will be stacked and padded. + - The output will contain 1 new dimension (dimension index 0) which will be the size of + of the original list. + + Args: + tensors (tp.List[torch.Tensor]): List of tensors to collate. + dim (int): Dimension which will be stacked and padded. + Returns: + tp.Tuple[torch.Tensor, torch.Tensor]: + torch.Tensor: Stacked and padded tensor. The output will contain 1 new dimension + (dimension index 0) which will be the size of the original list. + torch.Tensor: Tensor containing length of original tensor sizes (without padding). + """ + tensors = [x.transpose(0, dim) for x in tensors] + lens = torch.LongTensor([len(x) for x in tensors]) + padded_tensors = pad_sequence(tensors) + padded_tensors = padded_tensors.transpose(0, 1) + padded_tensors = padded_tensors.transpose(1, dim + 1) + return padded_tensors, lens + + +# TODO: Move to flashy? +def copy_state(state: tp.Any, device: tp.Union[torch.device, str] = 'cpu', + dtype: tp.Optional[torch.dtype] = None) -> tp.Any: + if isinstance(state, torch.Tensor): + if dtype is None or not state.is_floating_point(): + dtype = state.dtype + return state.detach().to(device=device, dtype=dtype, copy=True) + elif isinstance(state, dict): + return {k: copy_state(v, device, dtype) for k, v in state.items()} + elif isinstance(state, list): + return [copy_state(v, device, dtype) for v in state] + + +# TODO: Move to flashy? +@contextmanager +def swap_state(model, state, **kwargs): + old_state = copy_state(model.state_dict()) + model.load_state_dict(state, **kwargs) + try: + yield + finally: + model.load_state_dict(old_state) + + +@lru_cache(None) +def warn_once(logger, msg): + """Warn about a given message only once.""" + logger.warning(msg) + + +def is_jsonable(x: tp.Any): + """Check if an object can be serialized into a json:""" + try: + json.dumps(x) + return True + except (TypeError, OverflowError): + return False + + +def load_clap_state_dict(clap_model, path: tp.Union[str, Path]): + """Wrapper around state dict loading of CLAP model + addressing compatibility issues between CLAP and AudioCraft + HuggingFace transformer version. + See: https://github.com/LAION-AI/CLAP/issues/118 + """ + from clap_module.factory import load_state_dict # type: ignore + pkg = load_state_dict(path) + pkg.pop('text_branch.embeddings.position_ids', None) + clap_model.model.load_state_dict(pkg) diff --git a/src/lm_loopgen.py b/src/lm_loopgen.py new file mode 100644 index 0000000000000000000000000000000000000000..5c22f031ad2c3950e6837a36947ce7ec47504d06 --- /dev/null +++ b/src/lm_loopgen.py @@ -0,0 +1,375 @@ +# 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 logging +import math +import typing as tp +import torch +import tqdm + +from copy import deepcopy +from src.audiocraft.utils import utils +from src.audiocraft.modules.conditioners import ( + ConditioningAttributes, + ConditionType, +) +from src.audiocraft.models.lm import LMModel +from src.audiocraft.models.lm_magnet import MagnetLMModel + +logger = logging.getLogger(__name__) +ConditionTensors = tp.Dict[str, ConditionType] +CFGConditions = tp.Union[ConditionTensors, tp.Tuple[ConditionTensors, ConditionTensors]] + + +class LoopgenLMModel(MagnetLMModel): + + def __init__(self, subcodes_context: int = 5, compression_model_framerate: int = 50, + segment_duration: int = 10, span_len: int = 3, **kwargs): + super().__init__(subcodes_context, compression_model_framerate, + segment_duration, span_len, **kwargs) + + @torch.no_grad() + def generate(self, + valid_tokens: torch.Tensor, + left_tokens: tp.Optional[torch.Tensor] = None, + conditions: tp.List[ConditioningAttributes] = [], + negative_text_prompt: tp.Optional[tp.List[tp.Optional[str]]] = None, + num_samples: tp.Optional[int] = None, + max_gen_len: int = 256, + use_sampling: bool = True, + temp: float = 1.0, + top_k: int = 250, + top_p: float = 0.0, + cfg_coef: tp.Optional[float] = None, + cfg_coef_beta: tp.Optional[float] = None, + two_step_cfg: tp.Optional[bool] = None, + remove_prompts: bool = False, + check: bool = False, + callback: tp.Optional[tp.Callable[[int, int], None]] = None, + **kwargs) -> torch.Tensor: + + assert cfg_coef is None, "Unsupported in MAGNeT. Use max_cfg_coef,min_cfg_coef instead." + assert two_step_cfg is None, "MAGNeT currently doesn't support two step classifier-free-guidance." + assert remove_prompts is False, "MAGNeT currently doesn't support the remove_prompts arg." + assert check is False, "MAGNeT currently doesn't support the check arg." + assert cfg_coef_beta is None, "MAGNeT currently doesn't support the cfg_coef_beta arg." + # Call the MAGNeT-specific generation method + return self._generate_magnet( + valid_tokens=valid_tokens, + left_tokens=left_tokens, + conditions=conditions, + negative_text_prompt=negative_text_prompt, + num_samples=num_samples, + max_gen_len=max_gen_len, + use_sampling=use_sampling, + temp=temp, + top_k=top_k, + top_p=top_p, + callback=callback, **kwargs) + + @torch.no_grad() + def _generate_magnet(self, + valid_tokens: torch.Tensor, + left_tokens: tp.Optional[torch.Tensor] = None, + conditions: tp.List[ConditioningAttributes] = [], + negative_text_prompt: tp.Optional[tp.List[tp.Optional[str]]] = None, + num_samples: tp.Optional[int] = None, + max_gen_len: int = 256, + use_sampling: bool = True, + temp: float = 3.0, + top_k: int = 0, + top_p: float = 0.9, + callback: tp.Optional[tp.Callable[[int, int], None]] = None, + max_cfg_coef: float = 10.0, + min_cfg_coef: float = 1.0, + decoding_steps: tp.List[int] = [20, 10, 10, 10], + anneal_temp: bool = True, + rescorer: LMModel = None, + rescore_weights: torch.Tensor | float = 0.7, + rescorer_temp: torch.Tensor | float = 1.0, + offset: bool = True + ) -> torch.Tensor: + """Generate audio tokens given textual conditions, and optionally given audio prompts, + by running MAGNeT's iterative decoding algorithm for each of the n_q RVQ levels. + Args: + prompt (torch.Tensor): Prompt tokens of shape [B, K, T]. + conditions (list of ConditioningAttributes): List of conditions. + num_samples (int): Number of samples to generate when no prompt and no conditions are given. + max_gen_len (int): Maximum generation length. + use_sampling (bool): Whether to use a sampling strategy or not. + temp (float): Initial sampling temperature. + top_k (int): k for "top-k" sampling. + top_p (float): p for "top-p" sampling. + callback (Callback): Callback function to report generation progress. + max_clsfg_coef (float): Initial coefficient used for classifier free guidance. + min_clsfg_coef (float): Final coefficient used for classifier free guidance. + decoding_steps (list of n_q ints): The number of iterative decoding steps, + for each of the n_q RVQ codebooks. + anneal_temp (bool): When set to True, softmax temperature will be linearly decayed to zero, at each stage. + Returns: + torch.Tensor: Generated tokens. + """ + assert not self.training, "generation shouldn't be used in training mode." + first_param = next(iter(self.parameters())) + device = first_param.device + + # Checking all input shapes are consistent. + possible_num_samples = [] + if num_samples is not None: + possible_num_samples.append(num_samples) + elif left_tokens is not None: + possible_num_samples.append(len(left_tokens)) + elif conditions: + possible_num_samples.append(len(conditions)) + else: + possible_num_samples.append(1) + assert [x == possible_num_samples[0] for x in possible_num_samples], "Inconsistent inputs shapes" + num_samples = possible_num_samples[0] + + B, K = num_samples, self.num_codebooks + + mask_id = self.special_token_id + + # we generate codes with a fixed sequence length + shape = (B, K, max_gen_len) + + gen_codes = torch.full(shape, mask_id, dtype=torch.long, device=device) + + # filling the gen_codes with the prompt if needed + left_pads = torch.zeros(B, dtype=torch.long) + right_pads = torch.zeros(B, dtype=torch.long) + for i, valid_amount in enumerate(valid_tokens): + empty_amount = max_gen_len - valid_amount + left_pads[i] = min(valid_amount, empty_amount // 2) if offset else 0 + right_pads[i] = empty_amount - left_pads[i] + + rescorer_conditions = None + if rescorer is not None: + assert rescorer.special_token_id == mask_id, "Rescorer and generator should have the same mask id." + rescorer_conditions = rescorer.cfg_dropout(conditions) #* (loop_trick_rotations + 1)) + rescorer_conditions = rescorer.att_dropout(rescorer_conditions) + rescorer_conditions = rescorer.condition_provider.tokenize(rescorer_conditions) + # encode conditions and fuse, both have a streaming cache to not recompute when generating. + rescorer_conditions = rescorer.condition_provider(rescorer_conditions) + + if left_tokens is not None: + for i, tokens in enumerate(left_tokens): + gen_codes[i,..., left_pads[i]:left_pads[i]+tokens.shape[-1]] = tokens + + + # create the gen_sequence with proper interleaving from the pattern: [B, K, S] + gen_sequence = gen_codes + + + # below we create set of conditions: one conditional and one unconditional + # to do that we merge the regular condition together with the null condition + # we then do 1 forward pass instead of 2. + cfg_conditions: tp.Optional[ConditionTensors] + if conditions: + null_conditions = deepcopy(conditions) + for sample, negative_prompt in zip(null_conditions, negative_text_prompt): + for condition in sample.attributes["text"]: + sample.text[condition] = negative_prompt + + tokenized = self.condition_provider.tokenize(conditions + null_conditions) + cfg_conditions = self.condition_provider(tokenized) + else: + cfg_conditions = {} + + curr_step = 0 + pbar = tqdm.tqdm(total=sum(decoding_steps), desc="Generating", leave=False) + for stage, n_steps in zip(range(self.n_q), decoding_steps): + gen_sequence, curr_step = self._generate_stage(gen_sequence, + cfg_conditions, + stage=stage, + device=device, + valid_tokens=valid_tokens, + left_tokens=left_tokens, + temp=temp, + max_cfg_coef=max_cfg_coef, + min_cfg_coef=min_cfg_coef, + top_k=top_k, + top_p=top_p, + timesteps=n_steps, + anneal_temp=anneal_temp, + use_sampling=use_sampling, + curr_step=curr_step, + total_steps=sum(decoding_steps), + callback=callback, + rescorer=rescorer, + rescore_weights=rescore_weights, + rescorer_temp=rescorer_temp, + rescorer_conditions=rescorer_conditions, + offset=offset, + pbar=pbar) + + return gen_sequence + + @torch.no_grad() + def _generate_stage(self, + gen_sequence: torch.Tensor, + condition_tensors: tp.Optional[ConditionTensors], + stage: int, + device: torch.device, + valid_tokens: torch.Tensor, + left_tokens: tp.Optional[torch.Tensor] = None, + use_sampling: bool = True, + temp: float = 3.0, + max_cfg_coef: float = 10.0, + min_cfg_coef: float = 1.0, + top_k: int = 0, + top_p: float = 0.0, + timesteps: int = 10, + anneal_temp: bool = True, + curr_step: int = 0, + total_steps: int = 0, + callback: tp.Optional[tp.Callable[[int, int], None]] = None, + rescorer: LMModel = None, + rescore_weights: torch.Tensor | float = 0.7, + rescorer_temp: torch.Tensor | float = 1.0, + rescorer_conditions = tp.Optional[ConditionTensors], + offset: bool = True, + pbar: tqdm.tqdm = None) -> tp.Tuple[torch.Tensor, int]: + """Generate audio tokens of a single RVQ level (stage), given the previously generated stages, + and the textual conditions. + Args: + gen_sequence (torch.Tensor): Previously generated tokens. + condition_tensors (tp.Optional[ConditionTensors]): pre-computed conditioning tensors. + stage (int): RVQ level to generate. + device (torch.device): device of the output tensor. + prompt_length (int): Temporal length of the audio prompt. + prompt (torch.Tensor): Prompt tokens of shape [B, K, T]. + use_sampling (bool): Whether to use a sampling strategy or not. + temp (float): Initial sampling temperature. + max_clsfg_coef (float): Initial coefficient used for classifier free guidance. + min_clsfg_coef (float): Final coefficient used for classifier free guidance. + top_k (int): k for "top-k" sampling. + top_p (float): p for "top-p" sampling. + timesteps (int): Number of iterative decoding steps. + anneal_temp (bool): When set to True, softmax temperature will be linearly decayed to zero, at each stage. + curr_step (int): Global iterative decoding step counter. + total_steps (int): Total decoding steps. + callback (Callback): Callback function to report generation progress. + Returns: + tuple(torch.Tensor, int): Generated tokens and the current decoding step counter. + """ + B, K, T = gen_sequence.shape + shape = (B, 1, T) # generating a single codebook per stage + + mask_id = self.special_token_id + stage_gen_seq = torch.full(shape, mask_id, dtype=torch.long, device=device) + + DONT_REMASK_ME_SCORE = -1e4 + + model = self if self._fsdp is None else self._fsdp + + # token-wise scores + scores = torch.zeros(shape, dtype=torch.float32, device=device) + # scores[..., :prompt_length] = DONT_REMASK_ME_SCORE + # gen_T = T - prompt_length + gen_T = torch.ones(B) * T + + left_pads = torch.zeros(B, dtype=torch.long) + right_pads = torch.zeros(B, dtype=torch.long) + for i, valid_amount in enumerate(valid_tokens): + empty_amount = scores.shape[-1] - valid_amount + left_pads[i] = min(valid_amount, empty_amount // 2) if offset else 0 + right_pads[i] = empty_amount - left_pads[i] + scores[i, ..., :left_pads[i]] = DONT_REMASK_ME_SCORE + if right_pads[i] > 0: + scores[i, ..., -right_pads[i]:] = DONT_REMASK_ME_SCORE + gen_T[i] -= empty_amount + + if left_tokens is not None: + for i, tokens in enumerate(left_tokens): + scores[i,..., left_pads[i]:left_pads[i]+tokens.shape[-1]] = DONT_REMASK_ME_SCORE + stage_gen_seq[i,..., left_pads[i]:left_pads[i]+tokens.shape[-1]] = tokens[[stage]] + gen_T[i] -= tokens.shape[-1] + + if isinstance(rescore_weights, float): + rescore_weights = torch.ones(timesteps, device=device) * rescore_weights + + if isinstance(rescorer_temp, float): + rescorer_temp = torch.ones(timesteps, device=device) * rescorer_temp + + # run MAGNeT iterative decoding for "timesteps" iterations + for timestep, steps_left in zip(torch.linspace(0, 1, timesteps, device=device), reversed(range(timesteps))): + mask_p = torch.cos(timestep * math.pi * 0.5) + + # masking of the k least probable overlapping (stride 1) spans + mask = torch.concat(( + [self._least_probable_span_masking(scores[[i], :, :], max(int((mask_p * gen_T[i]).item()), 1)).to(device) + for i in range(B)]), dim=0) + stage_gen_seq[mask] = mask_id + + for i, (r, l) in enumerate(zip(right_pads, left_pads)): + if r > 0: + stage_gen_seq[i,..., :l] = stage_gen_seq[i,..., -l-r:-r] + stage_gen_seq[i,..., -r:] = torch.cat([stage_gen_seq[i,..., l:-r]] * (r // (T - r - l) + 1), -1)[..., :r] + + gen_sequence[:, [stage], :] = stage_gen_seq + sequence = gen_sequence + + if condition_tensors: + sequence = torch.cat([sequence, sequence], dim=0) + + all_logits = model(sequence, [], condition_tensors, stage=stage) + + if condition_tensors: + # classifier free guidance with annealing + cond_logits, uncond_logits = all_logits.split(all_logits.shape[0]//2, dim=0) # [B, K, T, card] + clsfg_coef = float(mask_p) * max_cfg_coef + (1 - float(mask_p)) * min_cfg_coef + logits = uncond_logits + (cond_logits - uncond_logits) * clsfg_coef + else: + logits = all_logits + + # temperature annealing - linear + t = temp * (steps_left / timesteps) if anneal_temp else temp + + # sampling + logits = logits[:, stage, :, :].unsqueeze(1) + probs = torch.softmax(logits / max(t, 1e-2), dim=-1) + + if use_sampling: + if top_p > 0.0: + sampled_tokens = utils.sample_top_p(probs, p=top_p) + elif top_k > 0: + sampled_tokens = utils.sample_top_k(probs, k=top_k) + else: + sampled_tokens = utils.multinomial(probs, num_samples=1) + else: + sampled_tokens = torch.argmax(logits, dim=-1, keepdim=True) + + # place mask_id token in each of the masked positions + mask = stage_gen_seq == mask_id + stage_gen_seq = torch.where(mask, sampled_tokens[..., 0], stage_gen_seq) + gen_sequence[:, [stage], :] = stage_gen_seq + + # get probs of sampled tokens + sampled_probs = torch.gather(probs, 3, sampled_tokens)[..., 0] + + if rescorer: + rescorer_logits = rescorer.compute_predictions(gen_sequence, conditions=None, condition_tensors=rescorer_conditions).logits[:, [stage]] + + rescorer_probs = torch.softmax(rescorer_logits / rescorer_temp[steps_left], dim=-1) + rescorer_sampled_probs = torch.gather(rescorer_probs, 3, sampled_tokens)[..., 0] + + sampled_probs = rescore_weights[[steps_left]] * rescorer_sampled_probs + (1 - rescore_weights[[steps_left]]) * sampled_probs + + # prod in log space for lps masking (stride1) + scores = -torch.log(sampled_probs) + + # Fix unmasked tokens by placing inf probs (-inf scores) + scores = scores.masked_fill(~mask, DONT_REMASK_ME_SCORE) + + if callback is not None: + curr_step += 1 + callback(curr_step, total_steps) + + if pbar is not None: + pbar.update(1) + + return gen_sequence, curr_step \ No newline at end of file diff --git a/src/loopgen.py b/src/loopgen.py new file mode 100644 index 0000000000000000000000000000000000000000..ad2ce6a6f0b6fd170b24a27bcbe5268d03a0c5fd --- /dev/null +++ b/src/loopgen.py @@ -0,0 +1,338 @@ +# 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. + +""" +Main model for using MAGNeT. This will combine all the required components +and provide easy access to the generation API. +""" +import typing as tp +import torch +import gc + +from src.audiocraft.models.magnet import MAGNeT +from src.audiocraft.models.lm import LMModel +from src.audiocraft.modules.conditioners import ConditioningAttributes +from src.audiocraft.data.audio_utils import convert_audio +from src.audiocraft.models.loaders import load_compression_model, load_lm_model_ckpt, _delete_param +from pathlib import Path +from omegaconf import OmegaConf +from src.audiocraft.models.builders import get_condition_fuser, get_conditioner_provider, get_codebooks_pattern_provider +from src.audiocraft.utils.utils import dict_from_config +from .lm_loopgen import LoopgenLMModel + +def load_lm_model_loopgen(file_or_url_or_id: tp.Union[Path, str], compression_model_frame_rate: int, + device='cpu', cache_dir: tp.Optional[str] = None): + pkg = load_lm_model_ckpt(file_or_url_or_id, cache_dir=cache_dir) + cfg = OmegaConf.create(pkg['xp.cfg']) + cfg.device = str(device) + if cfg.device == 'cpu': + cfg.dtype = 'float32' + else: + cfg.dtype = 'float16' + _delete_param(cfg, 'conditioners.args.merge_text_conditions_p') + _delete_param(cfg, 'conditioners.args.drop_desc_p') + + cfg.transformer_lm.compression_model_framerate = compression_model_frame_rate + cfg.transformer_lm.segment_duration = cfg.dataset.segment_duration + cfg.transformer_lm.span_len = cfg.masking.span_len + + # MAGNeT models v1 support only xformers backend. + from src.audiocraft.modules.transformer import set_efficient_attention_backend + + if cfg.transformer_lm.memory_efficient: + set_efficient_attention_backend("xformers") + + kwargs = dict_from_config(getattr(cfg, "transformer_lm")) + n_q = kwargs["n_q"] + q_modeling = kwargs.pop("q_modeling", None) + codebooks_pattern_cfg = getattr(cfg, "codebooks_pattern") + attribute_dropout = dict_from_config(getattr(cfg, "attribute_dropout")) + cls_free_guidance = dict_from_config(getattr(cfg, "classifier_free_guidance")) + cfg_prob, cfg_coef = ( + cls_free_guidance["training_dropout"], + cls_free_guidance["inference_coef"], + ) + fuser = get_condition_fuser(cfg) + condition_provider = get_conditioner_provider(kwargs["dim"], cfg).to(cfg.device) + if len(fuser.fuse2cond["cross"]) > 0: # enforce cross-att programmatically + kwargs["cross_attention"] = True + if codebooks_pattern_cfg.modeling is None: + assert ( + q_modeling is not None + ), "LM model should either have a codebook pattern defined or transformer_lm.q_modeling" + codebooks_pattern_cfg = OmegaConf.create( + {"modeling": q_modeling, "delay": {"delays": list(range(n_q))}} + ) + + pattern_provider = get_codebooks_pattern_provider(n_q, codebooks_pattern_cfg) + model = LoopgenLMModel( + pattern_provider=pattern_provider, + condition_provider=condition_provider, + fuser=fuser, + cfg_dropout=cfg_prob, + cfg_coef=cfg_coef, + attribute_dropout=attribute_dropout, + dtype=getattr(torch, cfg.dtype), + device=cfg.device, + **kwargs, + ).to(cfg.device) + model.load_state_dict(pkg['best_state']) + model.eval() + model.cfg = cfg + return model + + + +class LoopGen(MAGNeT): + """MAGNeT main model with convenient generation API. + Args: + See MusicGen class. + """ + def __init__(self, **kwargs): + super().__init__(**kwargs) + + @staticmethod + def get_pretrained(name: str = 'facebook/magnet-small-10secs', device=None): + """Return pretrained model, we provide six models: + - facebook/magnet-small-10secs (300M), text to music, 10-second audio samples. + # see: https://huggingface.co/facebook/magnet-small-10secs + - facebook/magnet-medium-10secs (1.5B), text to music, 10-second audio samples. + # see: https://huggingface.co/facebook/magnet-medium-10secs + - facebook/magnet-small-30secs (300M), text to music, 30-second audio samples. + # see: https://huggingface.co/facebook/magnet-small-30secs + - facebook/magnet-medium-30secs (1.5B), text to music, 30-second audio samples. + # see: https://huggingface.co/facebook/magnet-medium-30secs + - facebook/audio-magnet-small (300M), text to sound-effect (10-second samples). + # see: https://huggingface.co/facebook/audio-magnet-small + - facebook/audio-magnet-medium (1.5B), text to sound-effect (10-second samples). + # see: https://huggingface.co/facebook/audio-magnet-medium + """ + if device is None: + if torch.cuda.device_count(): + device = 'cuda' + else: + device = 'cpu' + + compression_model = load_compression_model(name, device=device) + lm = load_lm_model_loopgen(name, compression_model_frame_rate=int(compression_model.frame_rate), device=device) + + if 'self_wav' in lm.condition_provider.conditioners: + lm.condition_provider.conditioners['self_wav'].match_len_on_eval = True + + kwargs = {'name': name, 'compression_model': compression_model, 'lm': lm} + return LoopGen(**kwargs) + + def set_generation_params(self, use_sampling: bool = True, top_k: int = 0, + top_p: float = 0.9, temperature: float = 3.0, + max_cfg_coef: float = 10.0, min_cfg_coef: float = 1.0, + decoding_steps: tp.List[int] = [20, 10, 10, 10], + span_arrangement: str = 'nonoverlap', + rescorer: LMModel = None, + rescore_weights: torch.Tensor | float = 0.7, + rescorer_temp: torch.Tensor | float = 1.0, + offset: bool = True): + """Set the generation parameters for MAGNeT. + + Args: + use_sampling (bool, optional): Use sampling if True, else do argmax decoding. Defaults to True. + top_k (int, optional): top_k used for sampling. Defaults to 0. + top_p (float, optional): top_p used for sampling, when set to 0 top_k is used. Defaults to 0.9. + temperature (float, optional): Initial softmax temperature parameter. Defaults to 3.0. + max_cfg_coef (float, optional): Coefficient used for classifier free guidance. Defaults to 10.0. + min_cfg_coef (float, optional): End coefficient of classifier free guidance annealing. Defaults to 1.0. + decoding_steps (list of n_q ints, optional): The number of iterative decoding steps, for each of the n_q RVQ codebooks. + """ + if isinstance(rescore_weights, int): + rescore_weights = float(rescore_weights) + + self.generation_params = { + 'use_sampling': use_sampling, + 'temp': temperature, + 'top_k': top_k, + 'top_p': top_p, + 'max_cfg_coef': max_cfg_coef, + 'min_cfg_coef': min_cfg_coef, + 'decoding_steps': [int(s) for s in decoding_steps], + 'rescorer': rescorer, + 'rescore_weights': rescore_weights, + 'rescorer_temp': rescorer_temp, + 'offset' : offset + } + + def generate(self, text_prompt: tp.List[str], negative_text_prompt: tp.Optional[tp.List[tp.Optional[str]]] = None, progress: bool = False, valid_tokens = None) -> tp.Union[torch.Tensor, tp.Tuple[torch.Tensor, torch.Tensor]]: + """Generate samples conditioned on text. + + Args: + descriptions (list of str): A list of strings used as text conditioning. + progress (bool, optional): Flag to display progress of the generation process. Defaults to False. + """ + if type(text_prompt) is not list: + text_prompt = [text_prompt] + + text_prompt, prompt_tokens = self._prepare_tokens_and_attributes(text_prompt, None) + + if valid_tokens is None: + valid_tokens = torch.ones(len(text_prompt)) * int(self.duration * self.frame_rate) + + assert prompt_tokens is None + if type(negative_text_prompt) is not list: + negative_text_prompt = [negative_text_prompt] * len(text_prompt) + + tokens = self._generate_tokens(valid_tokens, text_prompt, negative_text_prompt, None, progress) + + if valid_tokens[0] == int(self.duration * self.frame_rate): + return self.generate_audio(tokens, valid_tokens, 1) + else: + return self.generate_audio(tokens, valid_tokens, 2) + + @torch.no_grad() + def _prepare_tokens_and_attributes( + self, + descriptions: tp.Sequence[tp.Optional[str]], + prompt: tp.Optional[torch.Tensor] + ) -> tp.Tuple[tp.List[ConditioningAttributes], tp.Optional[torch.Tensor]]: + """Prepare model inputs. + + Args: + descriptions (list of str): A list of strings used as text conditioning. + prompt (torch.Tensor): A batch of waveforms used for continuation. + """ + attributes = [ + ConditioningAttributes(text={'description': description}) + for description in descriptions] + + if prompt is not None: + prompt_tokens = [] + if descriptions is not None: + assert len(descriptions) == len(prompt), "Prompt and nb. descriptions doesn't match" + for p in prompt: + p, scale = self.compression_model.encode(p.to(self.device)) + assert scale is None + prompt_tokens.append(p[0]) + else: + prompt_tokens = None + + return attributes, prompt_tokens + + def generate_continuation(self, left_hint: torch.Tensor, left_hint_sr: int, text_prompt: tp.Optional[tp.List[tp.Optional[str]]] = None, negative_text_prompt: tp.Optional[tp.List[tp.Optional[str]]] = None, + progress: bool = False, valid_tokens = None) \ + -> tp.Union[torch.Tensor, tp.Tuple[torch.Tensor, torch.Tensor]]: + """Generate samples conditioned on audio prompts and an optional text description. + + Args: + prompt (torch.Tensor): A batch of waveforms used for continuation. + Prompt should be [B, C, T], or [C, T] if only one sample is generated. + prompt_sample_rate (int): Sampling rate of the given audio waveforms. + descriptions (list of str, optional): A list of strings used as text conditioning. Defaults to None. + progress (bool, optional): Flag to display progress of the generation process. Defaults to False. + """ + # if prompt.dim() == 2: + # prompt = prompt[None] + # if prompt.dim() != 3: + # raise ValueError("prompt should have 3 dimensions: [B, C, T] (C = 1).") + if type(text_prompt) is not list: + text_prompt = [text_prompt] + + if type(left_hint) is not list: + left_hint = [left_hint] + + audio_prompts = left_hint + left_hint = [] + for p in audio_prompts: + left_hint.append(convert_audio(p[None], left_hint_sr, self.sample_rate, self.audio_channels)) + + if text_prompt is None: + text_prompt = [None] * len(left_hint) + if valid_tokens is None: + valid_tokens = torch.ones(len(left_hint)) * int(self.duration * self.frame_rate) + + text_prompt, left_tokens = self._prepare_tokens_and_attributes(text_prompt, left_hint) + assert left_tokens is not None + + if type(negative_text_prompt) is not list: + negative_text_prompt = [negative_text_prompt] * len(text_prompt) + + tokens = self._generate_tokens(valid_tokens, text_prompt, negative_text_prompt, left_tokens, progress) + + if valid_tokens[0] == int(self.duration * self.frame_rate): + return self.generate_audio(tokens, valid_tokens, 1) + else: + return self.generate_audio(tokens, valid_tokens, 2) + + def _generate_tokens(self, valid_tokens: torch.Tensor, text_prompt: tp.List[ConditioningAttributes], negative_text_prompt: tp.Optional[tp.List[tp.Optional[str]]], + left_tokens: tp.Optional[torch.Tensor], progress: bool = False) -> torch.Tensor: + """Generate discrete audio tokens given audio prompt and/or conditions. + + Args: + attributes (list of ConditioningAttributes): Conditions used for generation (here text). + prompt_tokens (torch.Tensor, optional): Audio prompt used for continuation. + progress (bool, optional): Flag to display progress of the generation process. Defaults to False. + Returns: + torch.Tensor: Generated audio, of shape [B, C, T], T is defined by the generation params. + """ + total_gen_len = int(self.duration * self.frame_rate) + max_prompt_len = int(min(self.duration, self.max_duration) * self.frame_rate) + current_gen_offset: int = 0 + + def _progress_callback(generated_tokens: int, tokens_to_generate: int): + generated_tokens += current_gen_offset + if self._progress_callback is not None: + # Note that total_gen_len might be quite wrong depending on the + # codebook pattern used, but with delay it is almost accurate. + self._progress_callback(generated_tokens, tokens_to_generate) + else: + print(f'{generated_tokens: 6d} / {tokens_to_generate: 6d}', end='\r') + + tot_lr_tokens = 0 + if left_tokens is not None: + tot_lr_tokens +=left_tokens[0].shape[-1] + assert max_prompt_len >= tot_lr_tokens, \ + "Audio prompt is longer than available space" + + + callback = None + if progress: + callback = _progress_callback + + if self.duration <= self.max_duration: + # generate by sampling from LM, simple case. + with self.autocast: + gen_tokens = self.lm.generate( + valid_tokens, left_tokens, text_prompt, negative_text_prompt, + callback=callback, max_gen_len=total_gen_len, **self.generation_params) + + else: + assert False, "Prompt too long, we don't do that here" + return gen_tokens + + def generate_audio(self, gen_tokens: torch.Tensor, valid_tokens: torch.Tensor, k_loops: int) -> torch.Tensor: + """Generate Audio from tokens.""" + assert gen_tokens.dim() == 3 + with torch.no_grad(): + B, K, T = gen_tokens.shape + + offset = self.generation_params["offset"] + + left_pads = torch.zeros(B, dtype=torch.long) + right_pads = torch.zeros(B, dtype=torch.long) + for i, valid_amount in enumerate(valid_tokens): + empty_amount = T - valid_amount + left_pads[i] = min(valid_amount, empty_amount // 2) if offset else 0 + right_pads[i] = empty_amount - left_pads[i] + + gc.collect() + torch.cuda.empty_cache() + + gen_audio = self.compression_model.decode(gen_tokens, None) + + gen_audios = [] + for i, (r, l) in enumerate(zip(right_pads, left_pads)): + if r==0: + gen_audios.append(torch.cat([gen_audio[i,..., 640 * l:]] * k_loops, -1)) + else: + gen_audios.append(torch.cat([gen_audio[i,..., 640 * l: -640*r]] * k_loops, -1)) + + return gen_audios \ No newline at end of file