Module audiocraft.modules.jasco_conditioners

Classes

class ChordsEmbConditioner (card: int, out_dim: int, device: torch.device | str = 'cpu', **kwargs)
Expand source code 
class ChordsEmbConditioner(BaseConditioner):
    """
    A conditioner that embeds chord symbols into a continuous vector space.
    Attributes:
        card (int): The cardinality of the chord vocabulary.
        out_dim (int): The dimensionality of the output embeddings.
        device (Union[torch.device, str]): The device on which the embeddings are stored.
    """
    def __init__(self, card: int, out_dim: int, device: tp.Union[torch.device, str] = 'cpu', **kwargs):
        vocab_size = card + 1  # card + 1 - for null chord used during dropout
        super().__init__(dim=vocab_size, output_dim=-1)  # out_dim=-1 to avoid another projection
        self.emb = nn.Embedding(vocab_size, out_dim, device=device)
        self.device = device

    def tokenize(self, x: SymbolicCondition) -> SymbolicCondition:
        return SymbolicCondition(frame_chords=x.frame_chords.to(self.device))   # type: ignore

    def forward(self, x: SymbolicCondition) -> ConditionType:
        embeds = self.emb(x.frame_chords)
        mask = torch.ones_like(embeds[..., 0])
        return embeds, mask

A conditioner that embeds chord symbols into a continuous vector space.

Attributes

card : int
The cardinality of the chord vocabulary.
out_dim : int
The dimensionality of the output embeddings.
device : Union[torch.device, str]
The device on which the embeddings are stored.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

Class variables

var call_super_init : bool
var dump_patches : bool
var training : bool

Inherited members

class DrumsConditioner (out_dim: int,
sample_rate: int,
blurring_factor: int = 3,
cache_path: pathlib.Path | str | None = None,
compression_model_latent_dim: int = 128,
compression_model_framerate: float = 50,
segment_duration: float = 10.0,
device: torch.device | str = 'cpu',
**kwargs)
Expand source code 
class DrumsConditioner(WaveformConditioner):
    def __init__(self, out_dim: int, sample_rate: int, blurring_factor: int = 3,
                 cache_path: tp.Optional[tp.Union[str, Path]] = None,
                 compression_model_latent_dim: int = 128,
                 compression_model_framerate: float = 50,
                 segment_duration: float = 10.0,
                 device: tp.Union[torch.device, str] = 'cpu',
                 **kwargs):
        """Drum condition conditioner

        Args:
            out_dim (int): _description_
            sample_rate (int): _description_
            blurring_factor (int, optional): _description_. Defaults to 3.
            cache_path (tp.Optional[tp.Union[str, Path]], optional): path to precomputed cache. Defaults to None.
            compression_model_latent_dim (int, optional): latent dimensino. Defaults to 128.
            compression_model_framerate (float, optional): frame rate of the representation model. Defaults to 50.
            segment_duration (float, optional): duration in sec for each audio segment. Defaults to 10.0.
            device (tp.Union[torch.device, str], optional): device. Defaults to 'cpu'.
        """
        from demucs import pretrained
        self.sample_rate = sample_rate
        self.__dict__['demucs'] = pretrained.get_model('htdemucs').to(device)
        stem_sources: list = self.demucs.sources  # type: ignore
        self.stem_idx = stem_sources.index('drums')
        self.compression_model = None
        self.latent_dim = compression_model_latent_dim
        super().__init__(dim=self.latent_dim, output_dim=out_dim, device=device)
        self.autocast = TorchAutocast(enabled=device != 'cpu', device_type=self.device, dtype=torch.float32)
        self._use_masking = False
        self.blurring_factor = blurring_factor
        self.seq_len = int(segment_duration * compression_model_framerate)
        self.cache = None  # If you wish to train with EmbeddingCache, call self.create_embedding_cache(cache_path)

    def create_embedding_cache(self, cache_path):
        if cache_path is not None:
            self.cache = EmbeddingCache(Path(cache_path) / 'wav', self.device,
                                        compute_embed_fn=self._calc_coarse_drum_codes_for_cache,
                                        extract_embed_fn=self._load_drum_codes_chunk)

    @torch.no_grad()
    def _get_drums_stem(self, wav: torch.Tensor, sample_rate: int) -> torch.Tensor:
        """Get parts of the wav that holds the drums, 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)
            drum_stem = stems[:, self.stem_idx]  # extract relevant stems for drums conditioning
            return convert_audio(drum_stem, self.demucs.samplerate, self.sample_rate, 1)  # type: ignore

    def _temporal_blur(self, z: torch.Tensor):
        # z: (B, T, C)
        B, T, C = z.shape
        if T % self.blurring_factor != 0:
            # pad with reflect for T % self.temporal_blurring on the right in dim=1
            pad_val = self.blurring_factor - T % self.blurring_factor
            z = torch.nn.functional.pad(z, (0, 0, 0, pad_val), mode='reflect')
        z = z.reshape(B, -1, self.blurring_factor, C).sum(dim=2) / self.blurring_factor
        z = z.unsqueeze(2).repeat(1, 1, self.blurring_factor, 1).reshape(B, -1, C)
        z = z[:, :T]
        assert z.shape == (B, T, C)
        return z

    @torch.no_grad()
    def _extract_coarse_drum_codes(self, wav: torch.Tensor, sample_rate: int) -> torch.Tensor:
        assert self.compression_model is not None

        # stem separation of drums
        drums = self._get_drums_stem(wav, sample_rate)

        # continuous encoding with compression model
        latents = self.compression_model.model.encoder(drums)

        # quantization to coarsest codebook
        coarsest_quantizer = self.compression_model.model.quantizer.layers[0]
        drums = coarsest_quantizer.encode(latents).to(torch.int16)
        return drums

    @torch.no_grad()
    def _calc_coarse_drum_codes_for_cache(self, path: tp.Union[str, Path],
                                          x: WavCondition, idx: int,
                                          max_duration_to_process: float = 600) -> torch.Tensor:
        """Extract blurred drum latents 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)

        max_frames_to_process = int(max_duration_to_process * self.sample_rate)
        if wav.shape[-1] > max_frames_to_process:
            # process very long tracks in chunks
            start = 0
            codes = []
            while start < wav.shape[-1] - 1:
                wav_chunk = wav[..., start: start + max_frames_to_process]
                codes.append(self._extract_coarse_drum_codes(wav_chunk, self.sample_rate)[0])
                start += max_frames_to_process
            return torch.cat(codes)

        return self._extract_coarse_drum_codes(wav, self.sample_rate)[0]

    def _load_drum_codes_chunk(self, full_coarse_drum_codes: torch.Tensor, x: WavCondition, idx: int) -> torch.Tensor:
        """Extract a chunk of coarse drum codes from the full coarse drum codes 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 chunks from pre-computed drum codes.")
        assert self.compression_model is not None
        frame_rate = self.compression_model.frame_rate
        target_length = int(frame_rate * wav_length / self.sample_rate)
        target_length = max(target_length, self.seq_len)
        index = int(frame_rate * seek_time)
        out = full_coarse_drum_codes[index: index + target_length]
        # pad
        out = torch.cat((out, torch.zeros(target_length - out.shape[0], dtype=out.dtype, device=out.device)))
        return out.to(self.device)

    @torch.no_grad()
    def _get_wav_embedding(self, x: WavCondition) -> torch.Tensor:
        bs = x.wav.shape[0]
        if x.wav.shape[-1] <= 1:
            # null condition
            return torch.zeros((bs, self.seq_len, self.latent_dim), device=x.wav.device, dtype=x.wav.dtype)

        # extract coarse drum codes
        no_undefined_paths = all(p is not None for p in x.path)
        no_nullified_cond = x.wav.shape[-1] > 1
        if 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]
            codes = 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."
            codes = self._extract_coarse_drum_codes(x.wav, x.sample_rate[0])

        assert self.compression_model is not None
        # decode back to the continuous representation of compression model
        codes = codes.unsqueeze(1).permute(1, 0, 2)  # (B, T) -> (1, B, T)
        codes = codes.to(torch.int64)
        latents = self.compression_model.model.quantizer.decode(codes)

        latents = latents.permute(0, 2, 1)  # [B, C, T] -> [B, T, C]

        # temporal blurring
        return self._temporal_blur(latents)

    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

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.

Drum condition conditioner

Args

out_dim : int
description
sample_rate : int
description
blurring_factor : int, optional
description. Defaults to 3.
cache_path : tp.Optional[tp.Union[str, Path]], optional
path to precomputed cache. Defaults to None.
compression_model_latent_dim : int, optional
latent dimensino. Defaults to 128.
compression_model_framerate : float, optional
frame rate of the representation model. Defaults to 50.
segment_duration : float, optional
duration in sec for each audio segment. Defaults to 10.0.
device : tp.Union[torch.device, str], optional
device. Defaults to 'cpu'.

Ancestors

Class variables

var call_super_init : bool
var dump_patches : bool
var training : bool

Methods

def create_embedding_cache(self, cache_path)
Expand source code 
def create_embedding_cache(self, cache_path):
    if cache_path is not None:
        self.cache = EmbeddingCache(Path(cache_path) / 'wav', self.device,
                                    compute_embed_fn=self._calc_coarse_drum_codes_for_cache,
                                    extract_embed_fn=self._load_drum_codes_chunk)
def tokenize(self,
x: WavCondition) ‑> WavCondition
Expand source code 
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

Apply WavConditioner tokenization and populate cache if needed.

Inherited members

class JascoConditioningProvider (*args,
chords_card: int = 194,
sequence_length: int = 500,
melody_dim: int = 53,
**kwargs)
Expand source code 
class JascoConditioningProvider(ConditioningProvider):
    """
    A cond-provider that manages and tokenizes various types of conditioning attributes for Jasco models.
    Attributes:
        chords_card (int): The cardinality of the chord vocabulary.
        sequence_length (int): The length of the sequence for padding purposes.
        melody_dim (int): The dimensionality of the melody matrix.
    """
    def __init__(self, *args,
                 chords_card: int = 194,
                 sequence_length: int = 500,
                 melody_dim: int = 53, **kwargs):
        self.null_chord = chords_card
        self.sequence_len = sequence_length
        self.melody_dim = melody_dim
        super().__init__(*args, **kwargs)

    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)

        symbolic = self._collate_symbolic(inputs, self.conditioners.keys())

        assert set(text.keys() | wavs.keys() | symbolic.keys()).issubset(set(self.conditioners.keys())), (
            f"Got an unexpected attribute! Expected {self.conditioners.keys()}, ",
            f"got {text.keys(), wavs.keys(), symbolic.keys()}"
        )

        for attribute, batch in chain(text.items(), wavs.items(), symbolic.items()):
            output[attribute] = self.conditioners[attribute].tokenize(batch)
        return output

    def _collate_symbolic(self, samples: tp.List[ConditioningAttributes],
                          conditioner_keys: tp.Set) -> tp.Dict[str, SymbolicCondition]:
        output = {}

        # collate if symbolic cond exists
        if any(x in conditioner_keys for x in JascoCondConst.SYM.value):

            for s in samples:
                # hydrate with null chord if chords not exist - for inference support
                if (s.symbolic == {} or
                        s.symbolic[JascoCondConst.CRD.value].frame_chords is None or
                        s.symbolic[JascoCondConst.CRD.value].frame_chords.shape[-1] <= 1):  # type: ignore
                    # no chords conditioning - fill with null chord token
                    s.symbolic[JascoCondConst.CRD.value] = SymbolicCondition(
                        frame_chords=torch.ones(self.sequence_len, dtype=torch.int32) * self.null_chord)

                if (s.symbolic == {} or
                        s.symbolic[JascoCondConst.MLD.value].melody is None or
                        s.symbolic[JascoCondConst.MLD.value].melody.shape[-1] <= 1):  # type: ignore
                    # no chords conditioning - fill with null chord token
                    s.symbolic[JascoCondConst.MLD.value] = SymbolicCondition(
                        melody=torch.zeros((self.melody_dim, self.sequence_len)))

            if JascoCondConst.CRD.value in conditioner_keys:
                # pad to max
                max_seq_len = max(
                    [s.symbolic[JascoCondConst.CRD.value].frame_chords.shape[-1] for s in samples])  # type: ignore
                padded_chords = [
                    torch.cat((x.symbolic[JascoCondConst.CRD.value].frame_chords,   # type: ignore
                               torch.ones(max_seq_len -
                                          x.symbolic[JascoCondConst.CRD.value].frame_chords.shape[-1],  # type: ignore
                                          dtype=torch.int32) * self.null_chord))
                    for x in samples
                ]
                output[JascoCondConst.CRD.value] = SymbolicCondition(frame_chords=torch.stack(padded_chords))
            if JascoCondConst.MLD.value in conditioner_keys:
                melodies = torch.stack([x.symbolic[JascoCondConst.MLD.value].melody for x in samples])  # type: ignore
                output[JascoCondConst.MLD.value] = SymbolicCondition(melody=melodies)
        return output

A cond-provider that manages and tokenizes various types of conditioning attributes for Jasco models.

Attributes

chords_card : int
The cardinality of the chord vocabulary.
sequence_length : int
The length of the sequence for padding purposes.
melody_dim : int
The dimensionality of the melody matrix.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

Class variables

var call_super_init : bool
var dump_patches : bool
var training : bool

Inherited members

class MelodyConditioner (card: int, out_dim: int, device: torch.device | str = 'cpu', **kwargs)
Expand source code 
class MelodyConditioner(BaseConditioner):
    """
    A conditioner that handles melody conditioning from pre-computed salience matrix.
    Attributes:
        card (int): The cardinality of the melody matrix.
        out_dim (int): The dimensionality of the output projection.
        device (Union[torch.device, str]): The device on which the embeddings are stored.
    """
    def __init__(self, card: int, out_dim: int, device: tp.Union[torch.device, str] = 'cpu', **kwargs):
        super().__init__(dim=card, output_dim=out_dim)
        self.device = device

    def tokenize(self, x: SymbolicCondition) -> SymbolicCondition:
        return SymbolicCondition(melody=x.melody.to(self.device))  # type: ignore

    def forward(self, x: SymbolicCondition) -> ConditionType:
        embeds = self.output_proj(x.melody.permute(0, 2, 1))  # type: ignore
        mask = torch.ones_like(embeds[..., 0])
        return embeds, mask

A conditioner that handles melody conditioning from pre-computed salience matrix.

Attributes

card : int
The cardinality of the melody matrix.
out_dim : int
The dimensionality of the output projection.
device : Union[torch.device, str]
The device on which the embeddings are stored.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

Class variables

var call_super_init : bool
var dump_patches : bool
var training : bool

Inherited members