Class Pattern

stable_audio_tools/models/codebook_patterns.py:19–266

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

@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):