src/lm_loopgen.py

# 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