Module audiocraft.models.jasco

Main model for using JASCO. This will combine all the required components and provide easy access to the generation API.

Classes

class JASCO (chords_mapping_path='assets/chord_to_index_mapping.pkl', **kwargs)

Expand source code

class JASCO(BaseGenModel):
    """JASCO main model with convenient generation API.
    Args:
       chords_mapping_path: path to chords to index mapping pickle
       kwargs - See MusicGen class.
    """
    def __init__(self, chords_mapping_path='assets/chord_to_index_mapping.pkl', **kwargs):
        super().__init__(**kwargs)
        # JASCO operates over a fixed sequence length defined in it's config.
        self.duration = self.lm.cfg.dataset.segment_duration

        # load chord2index mapping of Chordino (https://github.com/ohollo/chord-extractor)
        assert os.path.exists(chords_mapping_path)
        self.chords_mapping = pickle.load(open(chords_mapping_path, "rb"))

        # set generation parameters
        self.set_generation_params()

    @staticmethod
    def get_pretrained(name: str = 'facebook/jasco-chords-drums-400M', device=None,
                       chords_mapping_path='assets/chord_to_index_mapping.pkl'):
        """Return pretrained model, we provide 2 models:
        1. facebook/jasco-chords-drums-400M: 10s music generation conditioned on
                                             text, chords and drums, 400M parameters.
        2. facebook/jasco-chords-drums-1B: 10s music generation conditioned on
                                           text, chords and drums, 1B parameters.
        """
        if device is None:
            if torch.cuda.device_count():
                device = 'cuda'
            else:
                device = 'cpu'

        compression_model = load_compression_model(name, device=device)
        lm = load_jasco_model(name, compression_model, device=device)

        kwargs = {'name': name,
                  'compression_model': compression_model,
                  'lm': lm,
                  'chords_mapping_path': chords_mapping_path}
        return JASCO(**kwargs)

    def set_generation_params(self,
                              cfg_coef_all: float = 5.0,
                              cfg_coef_txt: float = 0.0,
                              **kwargs):
        """Set the generation parameters for JASCO.

        Args:
            cfg_coef_all (float, optional): Coefficient used in multi-source classifier free guidance -
                                            all conditions term. Defaults to 5.0.
            cfg_coef_txt (float, optional): Coefficient used in multi-source classifier free guidance -
                                            text condition term. Defaults to 0.0.

        """
        self.generation_params = {
            'cfg_coef_all': cfg_coef_all,
            'cfg_coef_txt': cfg_coef_txt
        }
        self.generation_params.update(kwargs)

    def _unnormalized_latents(self, latents: torch.Tensor) -> torch.Tensor:
        """Unnormalize latents, shifting back to EnCodec's expected mean, std"""
        assert self.cfg is not None
        scaled = latents * self.cfg.compression_model_latent_std
        return scaled + self.cfg.compression_model_latent_mean

    def generate_audio(self, gen_latents: torch.Tensor) -> torch.Tensor:
        """Decode audio from generated latents"""
        assert gen_latents.dim() == 3  # [B, T, C]

        # unnormalize latents
        gen_latents = self._unnormalized_latents(gen_latents)
        return self.compression_model.model.decoder(gen_latents.permute(0, 2, 1))

    def _generate_tokens(self, attributes: tp.List[ConditioningAttributes],
                         prompt_tokens: tp.Optional[torch.Tensor], progress: bool = False) -> torch.Tensor:
        """Generate continuous audio latents given 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 latents, of shape [B, T, C].
        """
        total_gen_len = int(self.duration * self.frame_rate)
        max_prompt_len = int(min(self.duration, self.max_duration) * self.frame_rate)

        def _progress_callback(ode_steps: int, max_ode_steps: int):
            ode_steps += 1
            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(ode_steps, max_ode_steps)
            else:
                print(f'{ode_steps: 6d} / {max_ode_steps: 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

        # generate by sampling from the LM
        with self.autocast:
            total_gen_len = math.ceil(self.duration * self.compression_model.frame_rate)
            return self.lm.generate(
                   prompt_tokens, attributes,
                   callback=callback, max_gen_len=total_gen_len, **self.generation_params)

    def _prepare_chord_conditions(
            self,
            attributes: tp.List[ConditioningAttributes],
            chords: tp.Optional[tp.List[tp.Tuple[str, float]]],
    ) -> tp.List[ConditioningAttributes]:
        """
        Prepares chord conditions by translating symbolic chord progressions into a sequence of integers.
        This method updates the ConditioningAttributes with per-frame chords information.
        Args:
            attributes (List[ConditioningAttributes]):
                The initial attributes and optional tensor data.
            chords (List[Tuple[str, float]]):
                A list of tuples containing chord labels and their start times.
        Returns:
            List[ConditioningAttributes]:
                The updated attributes with frame chords integrated, alongside the original optional tensor data.
        """
        if chords is None or chords == []:
            for att in attributes:
                att.symbolic[JascoCondConst.CRD.value] = SymbolicCondition(frame_chords=-1 *
                                                                           torch.ones(1, dtype=torch.int32))
            return attributes

        # flip from (chord, start_time) to (start_time, chord)
        chords_time_first: tp.List[tuple[float, str]] = [(item[1], item[0]) for item in chords]

        # translate symbolic chord progression into a sequence of ints
        frame_chords = construct_frame_chords(min_timestamp=0,
                                              chord_changes=chords_time_first,
                                              mapping_dict=self.chords_mapping,
                                              prev_chord='',
                                              frame_rate=self.compression_model.frame_rate,
                                              segment_duration=self.duration)
        # update the attribute objects
        for att in attributes:
            att.symbolic[JascoCondConst.CRD.value] = SymbolicCondition(frame_chords=torch.tensor(frame_chords))
        return attributes

    @torch.no_grad()
    def _prepare_drums_conditions(self,
                                  attributes:
                                  tp.List[ConditioningAttributes],
                                  drums_wav: tp.Optional[torch.Tensor],
                                  ):
        # prepare drums cond
        for attr in attributes:
            if drums_wav is None:
                attr.wav[JascoCondConst.DRM.value] = WavCondition(
                    torch.zeros((1, 1, 1), device=self.device),
                    torch.tensor([0], device=self.device),
                    sample_rate=[self.sample_rate],
                    path=[None])
            else:
                if JascoCondConst.DRM.value not in self.lm.condition_provider.conditioners:
                    raise RuntimeError("This model doesn't support drums conditioning. ")

                expected_length = self.lm.cfg.dataset.segment_duration * self.sample_rate
                # trim if needed
                drums_wav = drums_wav[..., :expected_length]

                # pad if needed
                if drums_wav.shape[-1] < expected_length:
                    diff = expected_length - drums_wav.shape[-1]
                    diff_zeros = torch.zeros((drums_wav.shape[0], drums_wav.shape[1], diff),
                                             device=drums_wav.device, dtype=drums_wav.dtype)
                    drums_wav = torch.cat((drums_wav, diff_zeros), dim=-1)

                attr.wav[JascoCondConst.DRM.value] = WavCondition(
                    drums_wav.to(device=self.device),
                    torch.tensor([drums_wav.shape[-1]], device=self.device),
                    sample_rate=[self.sample_rate],
                    path=[None],
                )

        return attributes

    @torch.no_grad()
    def _prepare_melody_conditions(
            self,
            attributes: tp.List[ConditioningAttributes],
            melody: tp.Optional[torch.Tensor],
            expected_length: int,
            melody_bins: int = 53,
    ) -> tp.List[ConditioningAttributes]:
        """
        Prepares melody conditions by subtituting with pre-computed salience matrix.
        This method updates the ConditioningAttributes with per-frame chords information.
        Args:
            attributes (List[ConditioningAttributes]):
                The initial attributes and optional tensor data.
            chords (List[Tuple[str, float]]):
                A list of tuples containing chord labels and their start times.
        Returns:
            List[ConditioningAttributes]:
                The updated attributes with frame chords integrated, alongside the original optional tensor data.
        """
        for attr in attributes:
            if melody is None:
                melody = torch.zeros((melody_bins, expected_length))
            attr.symbolic[JascoCondConst.MLD.value] = SymbolicCondition(melody=melody)
        return attributes

    @torch.no_grad()
    def _prepare_temporal_conditions(
            self,
            attributes: tp.List[ConditioningAttributes],
            expected_length: int,
            chords: tp.Optional[tp.List[tp.Tuple[str, float]]],
            drums_wav: tp.Optional[torch.Tensor],
            salience_matrix: tp.Optional[torch.Tensor],
            melody_bins: int = 53,
    ) -> tp.List[ConditioningAttributes]:
        """
        Prepares temporal conditions (chords, drums).
        Args:
            attributes (List[ConditioningAttributes]): The initial attributes and optional tensor data.
            expected_length (int): The expected number of generated frames.
            chords (List[Tuple[str, float]]):  A list of tuples containing chord labels and their start times.
            drums_wav (List[Tuple[str, float]]): tensor of extracted drums wav.
            salience_matrix (List[Tuple[str, float]]): melody matrix.
            melody_bins (int): number of melody bins the model was trained with, only relevant if trained with melody.
        Returns:
            List[ConditioningAttributes]:
                The updated attributes after processing chord conditions.
        """
        attributes = self._prepare_chord_conditions(attributes=attributes, chords=chords)
        attributes = self._prepare_drums_conditions(attributes=attributes, drums_wav=drums_wav)
        attributes = self._prepare_melody_conditions(attributes=attributes, melody=salience_matrix,
                                                     expected_length=expected_length, melody_bins=melody_bins)
        return attributes

    @torch.no_grad()
    def generate_music(
        self, descriptions: tp.List[str],
        drums_wav: tp.Optional[torch.Tensor] = None,
        drums_sample_rate: int = 32000,
        chords: tp.Optional[tp.List[tp.Tuple[str, float]]] = None,
        melody_salience_matrix: tp.Optional[torch.Tensor] = None,
        iopaint_wav: tp.Optional[torch.Tensor] = None,
        segment_duration: float = 10.0,
        frame_rate: float = 50.0,
        melody_bins: int = 53,
        progress: bool = False, return_latents: bool = False) \
            -> tp.Union[torch.Tensor, tp.Tuple[torch.Tensor, torch.Tensor]]:
        """Generate samples conditioned on text and temporal conditions (chords, melody, drums).

        Args:
            descriptions (list of str): A list of strings used as text conditioning.
            chords (list of (str, float) tuples): Chord progression represented as chord, start time (sec), e.g.:
                                                  [("C", 0.0), ("F", 4.0), ("G", 6.0), ("C", 8.0)]
            melody_salience_matrix (torch.Tensor, optional): melody saliency matrix. Default=None.
            iopaint_wav (torch.Tensor, optional): in/out=painting waveform. Default=None.
            segment_duration (float): the segment duration the model was trained on. Default=None.
            frame_rate (float): the frame_rate model was trained on. Default=None.
            melody_bins (int): number of melody bins the model was trained with, only relevant if trained with melody.
            progress (bool, optional): Flag to display progress of the generation process. Defaults to False.
        """

        if drums_wav is not None:
            if drums_wav.dim() == 2:
                drums_wav = drums_wav[None]
            assert drums_wav.dim() == 3, "drums wav should have a shape [B, C, T]."
            drums_wav = convert_audio(drums_wav, drums_sample_rate, self.sample_rate, self.audio_channels)

        cond_attributes, prompt_tokens = self._prepare_tokens_and_attributes(descriptions=descriptions,
                                                                             prompt=None)

        # prepare temporal conds (symbolic / audio)
        jasco_attributes = self._prepare_temporal_conditions(attributes=cond_attributes,
                                                             expected_length=int(segment_duration * frame_rate),
                                                             chords=chords,
                                                             drums_wav=drums_wav,
                                                             salience_matrix=melody_salience_matrix,
                                                             melody_bins=melody_bins)
        assert prompt_tokens is None
        tokens = self._generate_tokens(jasco_attributes, prompt_tokens, progress)
        if return_latents:
            return self.generate_audio(tokens), tokens
        return self.generate_audio(tokens)

    @torch.no_grad()
    def generate(self, descriptions: tp.List[str], progress: bool = False, return_latents: 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.
        """
        return self.generate_music(descriptions=descriptions, progress=progress, return_latents=return_latents)

JASCO main model with convenient generation API.

Args

chords_mapping_path

path to chords to index mapping pickle

kwargs - See MusicGen class.

Ancestors

• BaseGenModel
• abc.ABC

Static methods

def get_pretrained(name: str = 'facebook/jasco-chords-drums-400M',
device=None,
chords_mapping_path='assets/chord_to_index_mapping.pkl')

Expand source code

@staticmethod
def get_pretrained(name: str = 'facebook/jasco-chords-drums-400M', device=None,
                   chords_mapping_path='assets/chord_to_index_mapping.pkl'):
    """Return pretrained model, we provide 2 models:
    1. facebook/jasco-chords-drums-400M: 10s music generation conditioned on
                                         text, chords and drums, 400M parameters.
    2. facebook/jasco-chords-drums-1B: 10s music generation conditioned on
                                       text, chords and drums, 1B parameters.
    """
    if device is None:
        if torch.cuda.device_count():
            device = 'cuda'
        else:
            device = 'cpu'

    compression_model = load_compression_model(name, device=device)
    lm = load_jasco_model(name, compression_model, device=device)

    kwargs = {'name': name,
              'compression_model': compression_model,
              'lm': lm,
              'chords_mapping_path': chords_mapping_path}
    return JASCO(**kwargs)

Return pretrained model, we provide 2 models: 1. facebook/jasco-chords-drums-400M: 10s music generation conditioned on text, chords and drums, 400M parameters. 2. facebook/jasco-chords-drums-1B: 10s music generation conditioned on text, chords and drums, 1B parameters.

Methods

def generate_audio(self, gen_latents: torch.Tensor) ‑> torch.Tensor

Expand source code

def generate_audio(self, gen_latents: torch.Tensor) -> torch.Tensor:
    """Decode audio from generated latents"""
    assert gen_latents.dim() == 3  # [B, T, C]

    # unnormalize latents
    gen_latents = self._unnormalized_latents(gen_latents)
    return self.compression_model.model.decoder(gen_latents.permute(0, 2, 1))

Decode audio from generated latents

def generate_music(self,
descriptions: List[str],
drums_wav: torch.Tensor | None = None,
drums_sample_rate: int = 32000,
chords: List[Tuple[str, float]] | None = None,
melody_salience_matrix: torch.Tensor | None = None,
iopaint_wav: torch.Tensor | None = None,
segment_duration: float = 10.0,
frame_rate: float = 50.0,
melody_bins: int = 53,
progress: bool = False,
return_latents: bool = False) ‑> torch.Tensor | Tuple[torch.Tensor, torch.Tensor]

Expand source code

@torch.no_grad()
def generate_music(
    self, descriptions: tp.List[str],
    drums_wav: tp.Optional[torch.Tensor] = None,
    drums_sample_rate: int = 32000,
    chords: tp.Optional[tp.List[tp.Tuple[str, float]]] = None,
    melody_salience_matrix: tp.Optional[torch.Tensor] = None,
    iopaint_wav: tp.Optional[torch.Tensor] = None,
    segment_duration: float = 10.0,
    frame_rate: float = 50.0,
    melody_bins: int = 53,
    progress: bool = False, return_latents: bool = False) \
        -> tp.Union[torch.Tensor, tp.Tuple[torch.Tensor, torch.Tensor]]:
    """Generate samples conditioned on text and temporal conditions (chords, melody, drums).

    Args:
        descriptions (list of str): A list of strings used as text conditioning.
        chords (list of (str, float) tuples): Chord progression represented as chord, start time (sec), e.g.:
                                              [("C", 0.0), ("F", 4.0), ("G", 6.0), ("C", 8.0)]
        melody_salience_matrix (torch.Tensor, optional): melody saliency matrix. Default=None.
        iopaint_wav (torch.Tensor, optional): in/out=painting waveform. Default=None.
        segment_duration (float): the segment duration the model was trained on. Default=None.
        frame_rate (float): the frame_rate model was trained on. Default=None.
        melody_bins (int): number of melody bins the model was trained with, only relevant if trained with melody.
        progress (bool, optional): Flag to display progress of the generation process. Defaults to False.
    """

    if drums_wav is not None:
        if drums_wav.dim() == 2:
            drums_wav = drums_wav[None]
        assert drums_wav.dim() == 3, "drums wav should have a shape [B, C, T]."
        drums_wav = convert_audio(drums_wav, drums_sample_rate, self.sample_rate, self.audio_channels)

    cond_attributes, prompt_tokens = self._prepare_tokens_and_attributes(descriptions=descriptions,
                                                                         prompt=None)

    # prepare temporal conds (symbolic / audio)
    jasco_attributes = self._prepare_temporal_conditions(attributes=cond_attributes,
                                                         expected_length=int(segment_duration * frame_rate),
                                                         chords=chords,
                                                         drums_wav=drums_wav,
                                                         salience_matrix=melody_salience_matrix,
                                                         melody_bins=melody_bins)
    assert prompt_tokens is None
    tokens = self._generate_tokens(jasco_attributes, prompt_tokens, progress)
    if return_latents:
        return self.generate_audio(tokens), tokens
    return self.generate_audio(tokens)

Generate samples conditioned on text and temporal conditions (chords, melody, drums).

Args

descriptions : list of str

A list of strings used as text conditioning.

chords (list of (str, float) tuples): Chord progression represented as chord, start time (sec), e.g.:

[("C", 0.0), ("F", 4.0), ("G", 6.0), ("C", 8.0)]

melody_salience_matrix : torch.Tensor, optional

melody saliency matrix. Default=None.

iopaint_wav : torch.Tensor, optional

in/out=painting waveform. Default=None.

segment_duration : float

the segment duration the model was trained on. Default=None.

frame_rate : float

the frame_rate model was trained on. Default=None.

melody_bins : int

number of melody bins the model was trained with, only relevant if trained with melody.

progress : bool, optional

Flag to display progress of the generation process. Defaults to False.

def set_generation_params(self, cfg_coef_all: float = 5.0, cfg_coef_txt: float = 0.0, **kwargs)

Expand source code

def set_generation_params(self,
                          cfg_coef_all: float = 5.0,
                          cfg_coef_txt: float = 0.0,
                          **kwargs):
    """Set the generation parameters for JASCO.

    Args:
        cfg_coef_all (float, optional): Coefficient used in multi-source classifier free guidance -
                                        all conditions term. Defaults to 5.0.
        cfg_coef_txt (float, optional): Coefficient used in multi-source classifier free guidance -
                                        text condition term. Defaults to 0.0.

    """
    self.generation_params = {
        'cfg_coef_all': cfg_coef_all,
        'cfg_coef_txt': cfg_coef_txt
    }
    self.generation_params.update(kwargs)

Set the generation parameters for JASCO.

Args

cfg_coef_all : float, optional

Coefficient used in multi-source classifier free guidance - all conditions term. Defaults to 5.0.

cfg_coef_txt : float, optional

Coefficient used in multi-source classifier free guidance - text condition term. Defaults to 0.0.

Inherited members

• BaseGenModel:
  • audio_channels
  • frame_rate
  • generate
  • generate_continuation
  • generate_unconditional
  • sample_rate
  • set_custom_progress_callback