sheetsage/representations/jukebox.py

import io
import logging
import warnings
from contextlib import redirect_stdout

import jukebox.hparams
import jukebox.make_models
import jukebox.utils.dist_utils
import librosa
import numpy as np
import torch

from ..utils import decode_audio, get_approximate_audio_length
from .base import Representation

_SAMPLE_RATE = 44100
_FRAME_HOP_SIZE = 128
_MIN_LENGTH_SAMPLES = (60 * _SAMPLE_RATE) + 16
_MAX_LENGTH_SAMPLES = (600 * _SAMPLE_RATE) - 96
_CHUNK_FRAMES = 8192
_CHUNK_SAMPLES = _CHUNK_FRAMES * _FRAME_HOP_SIZE

_SINGLETON = None


def init_jukebox_singleton(model="5b", num_layers=53, log=True):
    global _SINGLETON

    if _SINGLETON is None:
        # Set up device
        with redirect_stdout(io.StringIO()) as s:
            rank, local_rank, device = jukebox.utils.dist_utils.setup_dist_from_mpi()
            if log:
                logging.info(s.getvalue())

        # Set up hyperparams
        hps = jukebox.hparams.Hyperparams()
        hps.sr = _SAMPLE_RATE
        hps.n_samples = 3 if model == "5b_lyrics" else 8
        hps.name = "samples"
        chunk_size = 16 if model == "5b_lyrics" else 32
        max_batch_size = 3 if model == "5b_lyrics" else 16
        hps.levels = 3
        hps.hop_fraction = [0.5, 0.5, 0.125]

        # Load VQVAE
        vqvae, *priors = jukebox.make_models.MODELS[model]
        with redirect_stdout(io.StringIO()) as s:
            vqvae = jukebox.make_models.make_vqvae(
                jukebox.hparams.setup_hparams(
                    vqvae, dict(sample_length=_CHUNK_SAMPLES)
                ),
                device,
            )
            if log:
                logging.info(s.getvalue())

        # Set up language model
        if num_layers is not None:
            overrides = dict(prior_depth=num_layers)
        else:
            overrides = dict()
        with redirect_stdout(io.StringIO()) as s:
            lm = jukebox.make_models.make_prior(
                jukebox.hparams.setup_hparams(priors[-1], overrides), vqvae, device
            )
            if log:
                logging.info(s.getvalue())
        lm.prior.only_encode = True

        _SINGLETON = (model, num_layers, hps, vqvae, lm, device)
    else:
        if (model, num_layers) != _SINGLETON[:2]:
            raise Exception("Jukebox can only be initialized once")

    return _SINGLETON


class Jukebox(Representation):
    def __init__(self, num_layers=53, fp16=False, log=True):
        # NOTE: Layer 53 is the deepest that fit on a commodity 12GB card
        (
            _,
            _,
            self.hps,
            self.vqvae,
            self.lm,
            self.device,
        ) = init_jukebox_singleton(model="5b", num_layers=num_layers, log=log)
        self.fp16 = fp16

    @classmethod
    def decode_audio(cls, audio_path, offset=0.0, duration=None):
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            audio, sr = librosa.load(
                audio_path, sr=None, mono=False, offset=offset, duration=duration
            )
        if audio.ndim == 1:
            audio = audio[np.newaxis, :]
        audio = np.swapaxes(audio, 0, 1)
        audio = np.mean(audio, axis=1, keepdims=False)
        if sr != _SAMPLE_RATE:
            audio = librosa.resample(audio, orig_sr=sr, target_sr=_SAMPLE_RATE, res_type="kaiser_best")
        if audio.shape[0] > 0:
            norm_factor = np.abs(audio).max()
            if norm_factor > 0:
                audio /= norm_factor
        return audio

    def _codify_audio(
        self, audio, tqdm=lambda x: x, window_size=_CHUNK_SAMPLES, pad=True
    ):
        # NOTE: Ugly API for legacy test case.
        hop_size = _CHUNK_SAMPLES
        hop_size_frames = window_size // _FRAME_HOP_SIZE
        result = []
        for i in tqdm(list(range(0, audio.shape[0], hop_size))):
            context = audio[i : i + window_size]
            if pad and context.shape[0] < window_size:
                context = np.pad(context, (0, window_size - context.shape[0]))
            with torch.no_grad():
                context = torch.tensor(
                    context, dtype=torch.float32, device=self.device
                ).view(1, -1, 1)
                context_codified = self.vqvae.encode(context)[-1].view(-1).cpu().numpy()
            context_codified = context_codified[:hop_size_frames]
            result.append(context_codified)
        return np.concatenate(result, axis=0)

    def codify_audio(self, audio, tqdm=lambda x: x):
        return self._codify_audio(audio, tqdm=tqdm)

    def lm_activations(
        self,
        audio_codified,
        metadata_offset_seconds=0.0,
        metadata_total_length_seconds=None,
        metadata_artist=None,
        metadata_genre=None,
        metadata_lyrics=None,
        tqdm=lambda x: x,
    ):
        hop_size = _CHUNK_FRAMES
        window_size = _CHUNK_FRAMES
        if audio_codified.shape[0] % _CHUNK_FRAMES != 0:
            raise ValueError()

        # Compute metadata offset
        metadata_initial_offset = int(metadata_offset_seconds * _SAMPLE_RATE)
        metadata_initial_offset = (
            metadata_initial_offset // _FRAME_HOP_SIZE
        ) * _FRAME_HOP_SIZE
        assert metadata_initial_offset % _FRAME_HOP_SIZE == 0
        if metadata_initial_offset < 0:
            raise ValueError()

        # Compute metadata total length
        if metadata_total_length_seconds is None:
            metadata_total_length = audio_codified.shape[0] * _FRAME_HOP_SIZE
        else:
            metadata_total_length = int(metadata_total_length_seconds * _SAMPLE_RATE)
        metadata_total_length = max(metadata_total_length, _MIN_LENGTH_SAMPLES)
        metadata_total_length = min(metadata_total_length, _MAX_LENGTH_SAMPLES)
        metadata_total_length = (
            metadata_total_length // _FRAME_HOP_SIZE
        ) * _FRAME_HOP_SIZE
        assert metadata_total_length % _FRAME_HOP_SIZE == 0
        assert metadata_total_length >= _MIN_LENGTH_SAMPLES
        assert metadata_total_length <= _MAX_LENGTH_SAMPLES

        result = []
        for i in tqdm(list(range(0, audio_codified.shape[0], hop_size))):
            # Select context window
            context = audio_codified[i : i + window_size]
            metadata_offset = metadata_initial_offset + i * _FRAME_HOP_SIZE
            metadata_offset = min(
                metadata_offset,
                metadata_total_length - (context.shape[0] * _FRAME_HOP_SIZE),
            )
            metadata_offset = max(metadata_offset, 0)
            assert metadata_offset % _FRAME_HOP_SIZE == 0

            with torch.no_grad():
                # Context
                x = torch.tensor(context, dtype=torch.int64, device=self.device).view(
                    1, -1
                )

                # Conditioning info
                meta = dict(
                    artist="unknown" if metadata_artist is None else metadata_artist,
                    genre="unknown" if metadata_genre is None else metadata_genre,
                    total_length=metadata_total_length,
                    offset=metadata_offset,
                    lyrics="Placeholder lyrics which do not affect 5b"
                    if metadata_lyrics is None
                    else metadata_lyrics,
                )
                metas = [meta] * self.hps.n_samples
                labels = [None, None, self.lm.labeller.get_batch_labels(metas, "cuda")]
                x_cond, y_cond, _ = self.lm.get_cond(None, self.lm.get_y(labels[-1], 0))
                x_cond = x_cond[:1]
                y_cond = y_cond[:1]

                # Extract activations
                activations = (
                    self.lm.prior.forward(
                        x, x_cond=x_cond, y_cond=y_cond, fp16=self.fp16
                    )
                    .cpu()
                    .numpy()
                )
                if self.fp16:
                    activations = activations.astype(np.float16)
                result.append(activations[0])

                # Clear memory
                del x
                del labels
                del x_cond
                del y_cond
                torch.cuda.empty_cache()

        return np.concatenate(result, axis=0)

    def __call__(self, audio_path, offset=0.0, duration=None):
        audio = self.decode_audio(audio_path, offset=offset, duration=duration)
        if offset == 0.0 and duration is None:
            total_length = audio.shape[0] / _SAMPLE_RATE
        else:
            total_length = get_approximate_audio_length(audio_path)
        codified_audio = self.codify_audio(audio)
        activations = self.lm_activations(
            codified_audio,
            metadata_offset_seconds=offset,
            metadata_total_length_seconds=total_length,
        )
        activations = activations[: int(audio.shape[0] / _FRAME_HOP_SIZE)]
        rate = _SAMPLE_RATE / _FRAME_HOP_SIZE
        return rate, activations