src/models/components/model.py

import torch
from torch import nn
from .musicgen_cc import MusicGen
from .utilities.model_utils import freeze, print_trainable_parameters
import peft


def get_musicgen(sec, device):
    mg = MusicGen.get_pretrained(name='large', device=device)
    # org repo sets extend_stride=16 which causes huge offsets 
    # (extend_stride -> stride_tokens -> current_gen_offset)
    # sec - 1 keeps a 1s overlap between chunks for continuity.
    mg.set_generation_params(duration=sec, extend_stride=sec - 1, top_k=250)
    mg.lm.here()
    freeze(mg.lm)
    return mg


class CondMusicgen(nn.Module):
    def __init__(self, sec, device="cuda", top_k=250):
        super().__init__()
        mg = get_musicgen(sec, device)
        mg.generation_params["top_k"] = top_k
        self.musicgen = mg
        self.lm = mg.lm
        self.max_duration = sec
        self.frame_rate = 50


    def set_training(self):
        self.lm.train()

    def forward(self, input_code, text_description, embed_fn, num_samples=1, mode="train",
                total_gen_len=None, prompt_tokens=None):

        mg = self.musicgen
        lm = self.lm

        # attributes, _ = mg._prepare_tokens_and_attributes(text_description, None)

        if mode == "train":
            with mg.autocast:
                out = lm.compute_predictions(codes=input_code,
                                             embed_fn=embed_fn,
                                             conditions=text_description)
            return out
        elif mode == "inference":
            if total_gen_len is None:
                total_gen_len = int(mg.duration * mg.frame_rate)

            with mg.autocast:
                gen_tokens = lm.generate(embed_fn=embed_fn, num_samples=num_samples,
                                         prompt=None, conditions=text_description,
                                         callback=None, max_gen_len=total_gen_len,
                                         **mg.generation_params)
                return gen_tokens
        elif mode == "continuation":
            with mg.autocast:
                #if prompt_tokens is not None:
                #    print(prompt_tokens.shape)
                gen_tokens = lm.generate(embed_fn=embed_fn, num_samples=num_samples,
                                         prompt=prompt_tokens, conditions=text_description,
                                         callback=None, max_gen_len=total_gen_len, **mg.generation_params)
                return gen_tokens

    def generate(self, cp_fn,
                 text_description,
                 condition_audio_code,
                 num_samples):

        mg = self.musicgen
        lm = self.lm

        attributes, _ = mg._prepare_tokens_and_attributes(text_description, None)

        all_tokens = []
        stride_tokens = int(self.frame_rate * mg.extend_stride)
        current_gen_offset = 0
        prompt_length = 0
        prompt_tokens = None
        total_gen_len = condition_audio_code.shape[-1] - 1
        total_sec = total_gen_len / 50.
        while current_gen_offset + prompt_length < total_gen_len:
            time_offset = current_gen_offset / self.frame_rate
            chunk_duration = min(total_sec - time_offset, self.max_duration)
            max_gen_len = int(chunk_duration * self.frame_rate)
            if prompt_length >= max_gen_len:
                break
            #print("current_gen_offset / total ", current_gen_offset, "/", total_gen_len)
            with mg.autocast:
                condition_audio_code_clip = condition_audio_code[:, :, current_gen_offset:current_gen_offset + max_gen_len + 1]
                #print(cond_mask.shape, drums_clip.shape, piano_roll_clip.shape, chords_clip.shape, max_gen_len)
                embed_fn = cp_fn(condition_audio_code=condition_audio_code_clip,
                                 max_n_frames=max_gen_len,
                                 mode="inference")
                gen_tokens = lm.generate(num_samples=num_samples,
                                         embed_fn=embed_fn,
                                         prompt=prompt_tokens,
                                         conditions=attributes,
                                         callback=None, max_gen_len=max_gen_len, **mg.generation_params)
            if prompt_tokens is None:
                all_tokens.append(gen_tokens)
            else:
                all_tokens.append(gen_tokens[:, :, prompt_tokens.shape[-1]:])
            prompt_tokens = gen_tokens[:, :, stride_tokens:]
            prompt_length = prompt_tokens.shape[-1]
            current_gen_offset += stride_tokens
            if current_gen_offset > 50 * 80:
                break

        gen_tokens = torch.cat(all_tokens, dim=-1)
        return gen_tokens



class EmbFn:
    def __init__(self, activates, fn, start_layer, max_len, inference=False, skip=None):
        self.interval = None
        self.index = -1
        self.adaptor = None
        self.start_layer = start_layer
        self.activates = activates
        self.max_len = max_len
        self.fn = fn
        self.inference = inference
        self.skip = skip

    def get_adaptor(self, tag):
        index = self.index
        if index < self.start_layer or tag == "cross":
            return None, None
        i = index - self.start_layer
        adaptor, gate = self.fn(i, self.activates)
        # if self.adaptor is not None:
        #    adaptor = self.adaptor + adaptor
        return adaptor, gate



    def clear_state(self):
        self.qkv = {}
        torch.cuda.empty_cache()

    def crop(self, tag, x):

        if self.interval is not None:
            st, ed = self.interval
            if st >= self.max_len:
                st = self.max_len - 1
                ed = st + 1
            return x[:, :, st:ed, :]
        return x

    def get_cross_attention_src(self, src):
        return src

    def modify(self, x, dt_x, gate):
        # return dt_x * gate[:dt_x.shape[-2], :] + x
        return dt_x * gate + x






    def update_adaptor(self, adaptor):
        self.adaptor = adaptor

    def set_index(self, index):
        self.index = index

    def update_interval(self, st, ed):
        self.interval = [st, ed]


class CPTransformerLayer(nn.Module):
    def __init__(self, norm1, norm2, layer_scale_1, dropout1, self_attn, layer_scale_2,
                 autocast, linear1, linear2, activation, dropout, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.norm1 = norm1
        self.norm2 = norm2
        self.layer_scale_1 = layer_scale_1
        self.dropout1 = dropout1
        self.self_attn = self_attn
        self.layer_scale_2 = layer_scale_2
        self.autocast = autocast
        self.linear1 = linear1
        self.linear2 = linear2
        self.activation = activation
        self.dropout = dropout

    def _ff_block(self, x):
        return self.linear2(self.dropout(self.activation(self.linear1(x))))

    def forward(self, x, cond=None):
        with self.autocast:
            if cond is None:
                nx = self.norm1(x)
            else:
                nx = self.norm1(x) + cond
            q, k, v, o = self.self_attn(nx, nx, nx, emb_fn=None,
                                        attn_mask=None,
                                        key_padding_mask=None,
                                        need_weights=False, is_causal=False, return_qkv=True)
            x = x + self.layer_scale_1(self.dropout1(o))
            x = x + self.layer_scale_2(self._ff_block(self.norm2(x)))
        return q, k, v, x


class CPTransformer(nn.Module):
    def __init__(self, model, emb_fn, start_layer, latent_dim, autocast, stride=50 * 10):
        super().__init__()

        self.emb_fn = {
            "emb": emb_fn
        }

        new_layers = nn.ModuleList()

        hidden_dim = 2048
        cond_dim = latent_dim
        num_layers = len(model.layers) - start_layer
        max_n_frames = 500

        self.pos_emb = nn.Parameter(
            torch.randn(num_layers + 1, max_n_frames + 1, hidden_dim),
            requires_grad=True)
        # self.encodec_emb = nn.Linear(hidden_dim, latent_dim, bias=False)
        self.merge_linear = nn.ModuleList()
        # self.piano_roll_emb = nn.ModuleList()
        for i in range(start_layer, len(model.layers)):
            norm1 = model.layers[i].norm1
            norm2 = model.layers[i].norm2
            layer_scale_1 = model.layers[i].layer_scale_1
            dropout1 = model.layers[i].dropout1
            self_attn = model.layers[i].self_attn
            layer_scale_2 = model.layers[i].layer_scale_2
            linear1 = model.layers[i].linear1
            linear2 = model.layers[i].linear2
            activation = model.layers[i].activation
            dropout = model.layers[i].dropout
            new_layers.append(CPTransformerLayer(norm1=norm1,
                                                 norm2=norm2,
                                                 layer_scale_1=layer_scale_1,
                                                 dropout1=dropout1,
                                                 self_attn=self_attn,
                                                 linear1=linear1,
                                                 linear2=linear2,
                                                 activation=activation,
                                                 dropout=dropout,
                                                 layer_scale_2=layer_scale_2,
                                                 autocast=autocast))

            self.merge_linear.append(nn.Linear(cond_dim, hidden_dim, bias=False))
            # self.piano_roll_emb.append(nn.Linear(128, latent_dim, bias=False))

        self.layers = new_layers
        # self.gates = nn.Parameter(torch.zeros([num_layers, max_n_frames, 64]))
        self.gates = nn.Parameter(torch.zeros([num_layers]))
        freeze(self.layers)

        self.max_n_frames = max_n_frames
        self.start_layer = start_layer
        self.num_layers = num_layers
        self.stride = stride

    def fn(self, i, activates):
        if i >= self.num_layers:
            return None, None
        return activates[i]

    def forward(self, condition_audio_code, max_n_frames, mode, skip=None):
        max_n_frames = self.max_n_frames if max_n_frames is None else max_n_frames

        sum_code = sum([self.emb_fn["emb"][i](condition_audio_code[:, i]) for i in range(4)])


        condition_audio_code = sum_code
        # condition_audio_code = self.encodec_emb(sum_code)

        B, T, latent_dim = condition_audio_code.shape  # (batch_size, n_frames, latent_dim)

        o = self.pos_emb[0][None, :T].repeat(B, 1, 1)
        #print(o.shape, T)

        outs = []

        encoded_condition = condition_audio_code
        for i in range(len(self.layers)):
            # We conduct two pass transformer.
            # The first pass is to get multi-layer representation of the condition_audio_code.
            # The second pass is to fuse the condition back.

            # 1st pass
            # encoded_condition = encoded_condition + self.pos_emb[i + 1][None, :T].repeat(B, 1, 1)
            # _, _, _, encoded_condition = self.layers[i](x=encoded_condition, cond=None)

            # add positional encoding and send to the transformer
            embedding = self.merge_linear[i](encoded_condition) + self.pos_emb[i + 1][None, :T].repeat(B, 1, 1)
            # embedding = (encoded_condition) + self.pos_emb[i + 1][None, :T].repeat(B, 1, 1) + self.merge_linear[i](condition_audio_code)

            # 2nd pass
            q, k, v, o = self.layers[i](x=o, cond=embedding)  #
            if not mode == "train":
                outs.append([[torch.cat([q, q], 0),
                              torch.cat([k, k], 0),
                              torch.cat([v, v], 0)], self.gates[i]])
            else:
                outs.append([[q, k, v], self.gates[i]])

        emb_fn = EmbFn(activates=outs, fn=self.fn,
                       start_layer=self.start_layer,
                       max_len=max_n_frames,
                       inference=(mode == "inference"),
                       skip=skip)

        return emb_fn

    def save_weights(self, path):
        state_dict = {}
        sdict = self.state_dict()
        for n in sdict:
            if str.startswith(n, "layers"):
                continue
            state_dict[n] = sdict[n]
        torch.save(state_dict, path)

    def load_weights(self, path):
        self.load_state_dict(torch.load(path, map_location="cpu"), strict=False)


class Instructor(nn.Module):
    def __init__(self, sec, num_layers, latent_dim, top_k):
        '''The MusicGen model with instructor adapter.

        Args:
            sec: int, duration of the audio in seconds
            num_layers: int, number of layers of adapter in the transformer
            latent_dim: int, dimension of the latent space
        '''
        super().__init__()
        lm = CondMusicgen(sec, top_k=top_k)
        self.peft_model = lm
        self.musicgen = lm.musicgen
        self.cp_transformer = CPTransformer(self.musicgen.lm.transformer,
                                            emb_fn=self.musicgen.lm.emb,
                                            start_layer=48 - num_layers,
                                            latent_dim=latent_dim,
                                            autocast=self.musicgen.autocast)
        self.text_lora_config = peft.LoraConfig(target_modules=r".*\.cross_attention\.(q_proj|v_proj)",
                                                r=32,
                                                lora_alpha=64)

        self.peft_model.lm.transformer = peft.get_peft_model(self.peft_model.lm.transformer, self.text_lora_config)

    def set_training(self):
        self.peft_model.set_training()
        print_trainable_parameters(self)

    def save_weights(self, path):
        self.cp_transformer.save_weights(path)

    def load_weights(self, path):
        self.cp_transformer.load_weights(path)

    def forward(self, input_code, text_description, condition_audio_code,
                num_samples=8, mode="train", max_n_frames=None, prompt_tokens=None):

        if max_n_frames is None:
            max_n_frames = input_code.shape[-1]

        condition_audio_code = torch.cat([condition_audio_code, torch.ones_like(condition_audio_code[:,:, 0:1]) * 2048], dim=-1)

        embed_fn = self.cp_transformer.forward(condition_audio_code=condition_audio_code,
                                       max_n_frames=max_n_frames,
                                       mode=mode,
                                       skip=None)
        out = self.peft_model.forward(input_code,
                              text_description=text_description,
                              embed_fn=embed_fn,
                              mode=mode,
                              total_gen_len=max_n_frames,
                              prompt_tokens=prompt_tokens)
        return out

    def generate(self, text_description, condition_audio_code,
                 num_samples=1):
        out = self.peft_model.generate(cp_fn=self.cp_transformer,
                                       text_description=text_description,
                                       condition_audio_code=condition_audio_code,
                                       num_samples=num_samples,
                                       )
        return out