audiocraft/lm.py

import torch
from audiocraft.transformer import StreamingTransformer
from torch import nn
from transformers import T5EncoderModel, T5Tokenizer  # type: ignore

class T5(nn.Module):

    def __init__(self):

        super().__init__()
        self.output_proj = nn.Linear(1024,  # t5-large
                                     1536)  # lm hidden
        self.t5_tokenizer = T5Tokenizer.from_pretrained('t5-large', legacy=True)
        t5 = T5EncoderModel.from_pretrained('t5-large').train(mode=False)

        # this makes sure that the t5 is not part
        # of the saved checkpoint
        self.__dict__['t5'] = t5.to('cuda:0')

    def forward(self, prompt):
        with torch.set_grad_enabled(False), torch.autocast(device_type='cuda', dtype=torch.float32):

            bs = len(prompt) // 2
            d = self.t5_tokenizer(prompt,
                                    return_tensors='pt',
                                    padding=True).to(self.output_proj.bias.device)
            d['attention_mask'][bs:, :] = 0  # null condition t5 attn_mask should be zero

            x = self.t5(input_ids=d['input_ids'],
                            attention_mask=d['attention_mask']).last_hidden_state  # no kv
        # Float 16
        # > self.output_proj() is outside of autocast of t5 - however inside the autocast of lm thus computed in torch.float16
        x = self.output_proj(x)  # nn.Linear() - produces different result if there is no duplicate txt condition here
        x[bs:, :, :] = 0  # venv/../site-packages/audiocraft/modules/conditioners.py -> tokenize()
        return x


class LMModel(nn.Module):

    def __init__(self,
                 n_q = 4,
                 card = 2048,
                 dim = 1536
                 ):
        super().__init__()
        self.t5 = T5()
        self.card = card # 2048
        self.n_draw = 1  # draw > 1 tokens of different CFG scale
                         # batch size > 1 is slower from n_draw as calls transformer on larger batch
        self.emb = nn.ModuleList([nn.Embedding(self.card + 1, dim) for _ in range(n_q)])  # EMBEDDING HAS 2049
        self.transformer = StreamingTransformer()
        self.out_norm = nn.LayerNorm(dim, eps=1e-5)
        self.linears = nn.ModuleList([nn.Linear(dim, self.card, bias=False) for _ in range(n_q)])  # LINEAR DOESNT HAVE 2049

    def forward(self,
                sequence,
                condition_tensors=None,
                cache_position=None):

        bs, n_q, time_frames = sequence.shape # [bs, 4, time]

        input_ = sum([self.emb[k](sequence[:, k]) for k in range(n_q)])

        out = self.transformer(torch.cat([input_, input_], 0),  # duplicate null condition (bs x 2) for ClassifierFreeGuidance
                               cross_attention_src=condition_tensors,
                               cache_position=cache_position
                               )

        logits = torch.stack([self.linears[k](self.out_norm(out)) for k in range(n_q)], dim=1) # [2*bs, 4, 1,      2048]
        logits = 3 * logits[:bs, :, :, :] - self._scale * logits[bs:, :, :, :]                 # [  bs, 4, n_draw, 2048]

        k = 24
        logits = torch.softmax(logits / 1.0, dim=3)  # [bs, 4, 1, 2048]
        p, ix = torch.topk(logits, k, dim=3)  # p = [bs, 4, 1, 24], ix = [bs, 4, 1, 2048]
        # Exponential Distribution
        deflation = torch.empty_like(p).exponential_(lambd=1)
        p = p / deflation
        # divide large probs with exp(prob) If prob=.001 then 1/exp(1*.001) -> almost by 0  --> exp doesnt really produce (0, Inf)
        p = p.argmax(dim=3, keepdim=True)  # [bs, 4, n_draw, 24]
        tok = ix.gather(dim=3, index=p).to(torch.int64)  # [bs, 4, n_draw, 1]
        return tok[:, :, :, 0].transpose(1, 2)  # [bs, n_draw, 4]

    @torch.no_grad()
    def generate(self,
                 max_tokens=None,
                 text_condition=None):
        x = self.t5(text_condition)
        bs = x.shape[0] // 2  # has null conditions - bs*2*N_REPEAT applys in builders.py
        self._scale = .3 * torch.rand(1, 1, self.n_draw, 1, device=x.device) + 1.94
        cache_position = 0

        out_codes = torch.full((bs,
                                self.n_draw,
                                4,
                                4 + 3 + max_tokens),  # 4 + max_tokens + 4-1 to have sufficient to index the 1st antidiagonal of 4x4 + 4 xtra tokens
                               self.card,
                               dtype=torch.long,
                               device=x.device) # [bs, n_draw, 4, dur]

        # A/R
        for offset in range(0, max_tokens + 4 - 1):  # max_tokens + n_q - 1

            # extract diagonal via indexing out_codes[ [0, 1, 2, 3], [0, 1, 2, 3] ]
            next_token = self.forward(out_codes[:, 0, [0, 1, 2, 3], torch.tensor([3, 2, 1, 0]) + offset][:, :, None],  # index diagonal & exapnd to [bs, n_q, dur=1]
                                      #gen_sequence[:, 0, :, offset-1:offset],  # DIAGINDEXING for setting prediction of lm into gen_sequence THE GENSEQUENCE has to be un-delayed in the end [Because it has to be de-delayed for the vocoder then is actually only the lm input that requires to see the delay thus we could just feed by diaggather] so it matches gen_codes -1 a[[0, 1, 2, 3], torch.tensor([0, 1, 2, 3]) + 5]  the gen_sequence is indexed by vertical column and fed to lm however the prediction of lm is place diagonally with delay to the gen_sequence
                                      condition_tensors=x,  # utilisation of the attention mask of txt condition ?
                                      cache_position=cache_position)  # [bs, n_draw, 4]

            # Fill of next_token should be also placed on antidiagonal [not column]

            #   Do Not Overwrite 2048 of TRIU/TRIL = START/END => Do Not Fill them by Predicted Tokens
            # 0-th antidiagonal should be full of card = [2048, 2048, 2048, 2048]
            #
            #   [2048, 2048, 2048, 2048,    0,    1,    2,    3,    4,    5,    6, 2048, 2048, 2048],
            #   [2048, 2048, 2048, 2048, 2048,    0,    1,    2,    3,    4,    5,    6, 2048, 2048],
            #   [2048, 2048, 2048, 2048, 2048, 2048,    0,    1,    2,    3,    4,    5,    6, 2048],
            #   [2048, 2048, 2048, 2048, 2048, 2048, 2048,    0,    1,    2,    3,    4,    5,    6]]
            # NO OVerWriting
            if offset == 0:

                next_token[:, :, 1:4] = 2048  # self.card - bottom 3 entries of the antidiagonal should remain 2048

            elif offset == 1:

                next_token[:, :, 2:4] = 2048  # bottom 2 entries of the antidiagonal should remain 2048

            elif offset == 2:

                next_token[:, :, 3:4] = 2048

            elif offset == max_tokens:

                next_token[:, :, 0:1] = 2048  # top 1 entry of the antidiagonal should stay to 2048

            elif offset == (max_tokens + 1):

                next_token[:, :, 0:2] = 2048

            elif offset == (max_tokens + 2):

                next_token[:, :, 0:3] = 2048

            else:  # offset 3,4,5,6,7...... max_tokens-1   # FILL Complete n_q = 4 ANTIDIAGONAL ENTRIES

                pass #print('No delete anti-diag')

            out_codes[:, :, [0, 1, 2, 3], torch.tensor([3, 2, 1, 0]) + offset + 1] = next_token
            # Sink Attn
            if (offset > 0) and (offset % 71) == 0:
                n_preserve = 4
                self.transformer._flush(n_preserve=n_preserve)
                cache_position = n_preserve
            else:
                cache_position += 1

        # [bs, n_draw, 4, time+xtra] -> [bs, 4, n_draw, time] ->  [bs, 4, time * n_draw]
        out_codes = out_codes[:, :, :, 4:max_tokens+4].transpose(1, 2).reshape(bs, 4, self.n_draw * max_tokens)

        # flush for next API call
        self.transformer._flush()

        return out_codes  # SKIP THE 4 fill 2048