models/lightning_musicgen.py

import itertools
from pathlib import Path
import lightning as L
from typing import Dict, Any, List, Mapping, Type, Optional, Sequence
from torch import Tensor
import torch
from dataclasses import dataclass
from lightning.pytorch import utilities as lightning_utils
from tqdm import tqdm
import loralib as lora

import hyperparameters as hp
from conditioning.beat_embedder import Beat, SinusoidalBeatEmbedder
from conditioning.condition_dispatcher import ConditionDispatcher
from conditioning.condition_provider import ConditionProvider
from conditioning.condition_type import ConditionType
from conditioning.conditioning_method import ConditioningMethod
from conditioning.embedded_condition import EmbeddedCondition
from conditioning.t5embedder import T5Embedder, T5EmbedderCPU, T5EmbedderGPU
from models.encodec import EncodecModel
from conditioning.prompt_processor import (
    DefaultPromptProcessor, InterleavedContextPromptProcessor, PromptProcessor,
    StraightContextPromptProcessor)
from models.loss import compute_cross_entropy
from models.musicgen_lm import MusicgenLm
from utils.audio import load_audio, save_audio
from utils.inspection import sanity_check
from utils.sample import eval_decorator, sample_top_k
import config as cfg


class LightningMusicgen(L.LightningModule):

    def __init__(self, params: hp.MusicgenParams):
        super().__init__()
        self.params: hp.MusicgenParams = params

        # instantiate EnCodec
        self.encodec_model: EncodecModel = EncodecModel.from_params(
            self.params.encodec_params)
        self.sample_rate: int = self.encodec_model.sample_rate
        self.n_q: int = self.encodec_model.num_codebooks
        self.special_token: int = 2048
        assert self.n_q == 4

        # freeze encodec
        for p in self.encodec_model.parameters():
            p.requires_grad = False

        # instantiate prompt processor
        self.prompt_processor: PromptProcessor = self.params.prompt_processor_params.model_class(
            self.encodec_model,
            self.special_token,
            keep_only_valid_steps=params.prompt_processor_params.
            keep_only_valid_steps,
            context_dropout=params.prompt_processor_params.context_dropout)

        # check consistency in prompt processor and lm params
        if self.prompt_processor.uses_sep_token:
            if self.params.lm_params.sep_token is None:
                raise cfg.ConfigurationError(
                    "This prompt processor requires an LM with "
                    "support for a separator token.")

        # instantiate lm
        self.lm: MusicgenLm = MusicgenLm(self.params.lm_params)

        # check consistency in conditioning parameters
        if (self.params.conditioning_params.embedder_types.keys()
                != self.params.conditioning_params.conditioning_methods.keys()):
            t1 = set(
                p.value
                for p in self.params.conditioning_params.embedder_types.keys())
            t2 = set(
                p.value for p in
                self.params.conditioning_params.conditioning_methods.keys())
            raise ValueError(
                "Embeddings produced by the condition provider don't match the "
                f"conditioning methods given in params. "
                f"processed conditions: {t1} "
                f"conditioning methods: {t2}")

        # instantiate condition provider (embedder)
        self.condition_provider = ConditionProvider(
            self.params.lm_params.dim,
            embedder_types=self.params.conditioning_params.embedder_types,
        )

        if self.params.lm_params.weights is not None:
            w = torch.load(Path(self.params.lm_params.weights),
                           map_location=None,
                           weights_only=True)
            newstatedict = {
                "condition_provider.embedders.description.output_proj.weight":
                    w['conditioner.output_proj.weight'],
                "condition_provider.embedders.description.output_proj.bias":
                    w['conditioner.output_proj.bias']
            }
            self.load_state_dict(newstatedict, strict=False)

        # instantiate condition dispatcher (contains fusers)
        self.condition_dispatcher: ConditionDispatcher = ConditionDispatcher(
            self.params.conditioning_params.conditioning_methods,
            self.params.lm_params.dim,
            self.params.conditioning_params.conditioning_dropout)

        # inject lora if needed
        if self.params.lora_params is not None:
            self._inject_lora()

        self.save_hyperparameters()

    def _inject_lora(self):
        assert self.params.lora_params is not None

        # freeze the decoder model
        # for p in self.lm.decoder.parameters():
        #     p.requires_grad = False

        layers: torch.nn.ModuleList = self.lm.decoder.attn_layers.layers
        # inject lora in every attention layer
        for att_idx in range(self.params.lm_params.n_layers):
            for sublayer_idx in range(2):
                layeridx = att_idx * 3 + sublayer_idx
                sublayer: torch.nn.ModuleList = layers[layeridx]  # type: ignore
                # for all layer types we want to swap
                for layername in self.params.lora_params.layers:
                    source_layer = sublayer[1].__getattr__(f"to_{layername}")
                    new_layer = lora.Linear(source_layer.in_features,
                                            source_layer.out_features,
                                            self.params.lora_params.r,
                                            self.params.lora_params.alpha,
                                            self.params.lora_params.dropout,
                                            bias=False)
                    with torch.no_grad():
                        new_layer.weight.data.copy_(source_layer.weight.data)
                    if layername == "q":
                        sublayer[1].to_q = new_layer
                    elif layername == "k":
                        sublayer[1].to_k = new_layer
                    elif layername == "v":
                        sublayer[1].to_v = new_layer
                    elif layername == "out":
                        sublayer[1].to_out = new_layer
                    else:
                        raise RuntimeError(f"unknown layer name {layername}")
                    # sublayer[1].__setattr__(f"to_{layername}", new_layer)

        lora.mark_only_lora_as_trainable(self)

    def configure_optimizers(self):  # type: ignore
        opt = torch.optim.AdamW(self.parameters(),
                                lr=1e-5,
                                betas=(0.9, 0.95),
                                weight_decay=0.001)
        return opt
        # n_warmup_steps: int = 1000

        # projections = (
        #     p for n, p in self.named_parameters() if "output_proj" in n)
        # embeddings = (p for n, p in self.named_parameters() if "token_emb" in n)
        # warmup_params = itertools.chain(projections, embeddings)
        # others = (p for n, p in self.named_parameters()
        #           if "output_proj" not in n and "token_emb" not in n)

        # multigroup_optim = torch.optim.AdamW(
        #     ({
        #         "params": warmup_params
        #     }, {
        #         "params": others
        #     }),
        #     lr=2e-5,
        #     betas=(0.9, 0.95),
        #     weight_decay=0.1,
        # )

        # lambda_encoder = lambda x: (1 + 1.5 * (1 - (x / n_warmup_steps))
        #                            ) if x < n_warmup_steps else 1.
        # lambda_decoder = lambda x: 0. if x < n_warmup_steps else 1.

        # multigroup_scheduler = torch.optim.lr_scheduler.LambdaLR(
        #     multigroup_optim,
        #     lr_lambda=[lambda_encoder, lambda_decoder],
        # )

        # scheduler_config = {
        #     "scheduler": multigroup_scheduler,
        #     "interval": "step"
        # }

        # return {"optimizer": multigroup_optim, "lr_scheduler": scheduler_config}

    # def on_train_batch_end(self, outputs: Tensor | Mapping[str, Any] | None,
    #                        batch: Any, batch_idx: int) -> None:
    # def on_before_optimizer_step(self, optimizer):
    #     decoder_grads = lightning_utils.grad_norm(self.lm.decoder, 2)
    #     self.log_dict(decoder_grads)

    def training_step(self, batch, batch_idx) -> Tensor:
        self.train(True)
        loss = self.run_step(batch)
        self.log(
            "train/loss",
            loss,
            # prog_bar=True,
            batch_size=len(batch["target"]),
            # sync_dist=True,
            # on_step=True,
        )
        self.log("global_step", self.global_step, prog_bar=True, logger=False)
        if self._trainer is not None and self.lr_schedulers() is not None:
            self.log(
                "train/new_params_lr",
                self.lr_schedulers().get_last_lr()[0],  # type: ignore
                prog_bar=True)
            self.log(
                "train/old_params_lr",
                self.lr_schedulers().get_last_lr()[1],  # type: ignore
                prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx) -> Tensor:
        self.train(False)
        with torch.no_grad():
            loss = self.run_step(batch)
        self.log(
            "val/loss",
            loss,
            prog_bar=True,
            batch_size=len(batch["target"]),
            sync_dist=True,
        )
        # sanity_check(self, interrupt=True)
        return loss

    def run_step(self, batch: Dict[str, Any]) -> Tensor:
        """
            Expects batch to be a dictionary like: 
            {
                "target": Tensor,
                "context": Tensor,     - optional
                "style": Tensor,       - optional
                "description": string, - optional
            }
        """

        # call prompt pre-processor
        (prompt, prompt_mask, target,
         decode_logits_fn) = self.prompt_processor.preprocess(batch)

        attention_mask = prompt_mask.sum(dim=-2) > 0

        # embed/encode conditioning data
        processed_conditions: Dict[ConditionType, EmbeddedCondition] = (
            self.condition_provider.process_conditions(batch))

        # dispatch eatch conditioning to the proper method, fusing if necessary
        method_to_cond: Dict[ConditioningMethod,
                             EmbeddedCondition] = self.condition_dispatcher(
                                 processed_conditions)

        # call language model
        logits = self.lm(
            x=prompt,
            attention_mask=attention_mask,
            cross_attention_input=method_to_cond.get(
                ConditioningMethod.CROSS_ATTENTION),
            prepend_embeds=method_to_cond.get(ConditioningMethod.INPUT_PREPEND),
            sum_embeds=method_to_cond.get(ConditioningMethod.INPUT_SUM),
        )

        # de-interleave logits and postprocess prompt
        logits, logits_mask = decode_logits_fn(logits)

        # compute cross entropy
        cross_entropy_loss, _ = compute_cross_entropy(logits, target,
                                                      logits_mask)

        return cross_entropy_loss

    def sample_next_token(
            self, current_sequence: Tensor, attention_mask: Tensor,
            method_to_cond: Dict[ConditioningMethod,
                                 EmbeddedCondition]) -> Tensor:

        if not current_sequence.isfinite().all():
            if current_sequence.isnan().any():
                print(f"Before forward pass some logits are nan")
            else:
                print(f"Before forward pass some logits are not finite")

        # call language model
        logits: Tensor = self.lm(
            x=current_sequence,
            attention_mask=attention_mask,
            cross_attention_input=method_to_cond.get(
                ConditioningMethod.CROSS_ATTENTION),
            prepend_embeds=method_to_cond.get(ConditioningMethod.INPUT_PREPEND),
            sum_embeds=method_to_cond.get(ConditioningMethod.INPUT_SUM),
        )

        if not logits.isfinite().all():
            if logits.isnan().any():
                print(f"After forward pass some logits are nan")
            else:
                print(f"After forward pass some logits are not finite")

        # classifier-free guidance
        cond_logits, uncond_logits = logits.split(
            current_sequence.shape[0] // 2,
            dim=0,
        )
        logits = uncond_logits + (cond_logits - uncond_logits) * 3.0

        # get logits for last token
        logits = logits.permute(0, 1, 3, 2)  # B, K, card, T
        logits = logits[..., -1]  # B, K, card,

        # apply softmax
        probs = torch.softmax(logits, dim=-1)

        # sample
        next_token = sample_top_k(probs, k=250)
        return next_token

    def predict_step(self, batch):
        batch["prog_bar"] = False
        return self.generate(**batch)

    # @torch.inference_mode()
    @eval_decorator
    @torch.no_grad()
    def generate(self,
                 n_samples: int,
                 gen_seconds: float | int,
                 prompt: Optional[Tensor],
                 context: Optional[Tensor | List[Tensor]],
                 style: Optional[Tensor],
                 beat: Optional[List[Beat]],
                 description: Optional[List[str]],
                 context_dropout_mask: Optional[Tensor] = None,
                 prog_bar: bool = False) -> Tensor:
        """Run autoregressive generation

        Args:
            n_samples (int): number of samples to generate (batch size). All other input parameters should match this.
            gen_seconds (float | int): total length of generation in seconds, including prompt if present.
            prompt (Optional[Tensor]): a piece of input to continue.
            context (Optional[Tensor  |  List[Tensor]]): a musical context to generate an accompaniment for.
            style (Optional[Tensor]): a piece of music to use as stylistic reference.
            beat (Optional[List[Beat]]): a beat object to follow
            description (Optional[List[str]]): a list of descriptions to use as conditioning
            context_dropout_mask (Optional[Tensor], optional): Defaults to None.
            prog_bar (bool, optional): whether to display a progress bar. Defaults to False.

        Raises:
            ValueError: _description_

        Returns:
            Tensor: _description_
        """

        n_gen_frames = int(self.encodec_model.frame_rate * gen_seconds)

        # generate empty sequence
        gen_sequence = torch.full((n_samples, self.n_q, n_gen_frames),
                                  -1,
                                  dtype=torch.long,
                                  device=self.device)

        # pre-process prompt to feed to the lm
        (gen_sequence, gen_mask, start_offset,
         decode_sequence_fn) = self.prompt_processor.prepare_for_generation(
             prompt,
             context,
             gen_sequence,
             use_cfg=True,
             context_dropout_mask=context_dropout_mask)

        # attention mask: in timesteps in which ALL residual layers are invalid,
        # set attention mask to False.
        attention_mask = gen_mask.sum(dim=-2) > 0

        # from now on we only need the first part of the gen_mask, the second was cfg
        gen_mask = gen_mask[:gen_mask.shape[0] // 2]

        # embed/encode conditioning data
        conditions = {
            "description": description,
            # "context": context,
            "style": style,
            "beat": beat,
        }

        # check for compatibilty of conditions
        for c_name, c_value in conditions.items():
            if c_value is None:
                continue
            condtype: ConditionType = ConditionType(c_name)
            if condtype.value not in self.condition_provider.embedders:
                raise ValueError(
                    f"This version of the model does not support conditioning "
                    f"with {c_name}. You should pass None.")

        processed_conditions: Dict[ConditionType, EmbeddedCondition] = (
            self.condition_provider.process_conditions(conditions,
                                                       duplicate_for_cfg=True,
                                                       batch_size=n_samples))

        # dispatch eatch conditioning to the proper method, fusing if necessary
        method_to_cond: Dict[ConditioningMethod,
                             EmbeddedCondition] = self.condition_dispatcher(
                                 processed_conditions)
        # autoregression

        iterator = range(start_offset, gen_sequence.shape[-1])
        if prog_bar:
            iterator = tqdm(iterator, desc="generating autoregressively...")

        for offset in iterator:
            current_sequence = gen_sequence[..., :offset]
            current_mask = attention_mask[..., :offset]
            next_token = self.sample_next_token(current_sequence, current_mask,
                                                method_to_cond)
            valid_mask = gen_mask[
                ...,  # TODO: I can't figure out if this is correct or if it matters at all anyways
                offset:offset + 1].expand(n_samples, -1, -1)
            next_token[~valid_mask] = self.special_token
            gen_sequence[:n_samples, :, offset:offset + 1] = torch.where(
                gen_sequence[:n_samples, :, offset:offset + 1] == -1,
                next_token,
                gen_sequence[:n_samples, :, offset:offset + 1],
            )
            gen_sequence[n_samples:, :, offset:offset + 1] = torch.where(
                gen_sequence[n_samples:, :, offset:offset + 1] == -1,
                next_token,
                gen_sequence[n_samples:, :, offset:offset + 1],
            )

            if prog_bar and torch.cuda.is_available():
                torch.cuda.synchronize()

        assert not (gen_sequence == -1).any()
        gen_sequence = gen_sequence[:gen_sequence.shape[0] // 2]
        # assert (gen_sequence == torch.where(
        #     gen_mask[None, ...].expand(n_samples, -1, -1),
        #     gen_sequence,
        #     self.special_token,
        # )).all()

        out_codes, out_mask = decode_sequence_fn(gen_sequence)

        self.encodec_model.eval()
        with torch.no_grad():
            out_audio = self.encodec_model.decode(out_codes)
        return out_audio

    @staticmethod
    def load_from_checkpoint_replacing_paths(ckp_path: Path | str):
        ckp_path = Path(ckp_path)

        def swap_parent(filepath: Path | str, new_parent: Path) -> Path | str:
            if isinstance(filepath, Path):
                typeout = Path
            elif isinstance(filepath, str):
                typeout = str
            else:
                raise RuntimeError("expected Path or str")

            filepath = Path(filepath)
            return typeout(new_parent / filepath.name)

        ckp = torch.load(ckp_path, map_location="cpu")
        params = ckp["hyper_parameters"]["params"]
        params.encodec_params.weights = swap_parent(
            params.encodec_params.weights, cfg.weights_dir())
        params.lm_params.weights = swap_parent(params.lm_params.weights,
                                               cfg.weights_dir())
        model = LightningMusicgen(params)
        model.load_state_dict(ckp["state_dict"])
        return model


if __name__ == "__main__":
    from time import time

    device = torch.device("cuda")

    # musicgen params
    model_params = hp.MusicgenParams(
        encodec_params=hp.pretrained_encodec_meta_32khz_params,
        prompt_processor_params=hp.PromptProcessorParams(
            model_class=InterleavedContextPromptProcessor,
            keep_only_valid_steps=True,
            context_dropout=0.5,
        ),
        conditioning_params=hp.ConditioningParams(
            embedder_types={
                ConditionType.DESCRIPTION: T5EmbedderGPU,
                ConditionType.BEAT: SinusoidalBeatEmbedder,
            },
            conditioning_methods={
                ConditionType.DESCRIPTION: ConditioningMethod.CROSS_ATTENTION,
                ConditionType.BEAT: ConditioningMethod.INPUT_PREPEND,
            },
            conditioning_dropout=0.5,
        ),
        lm_params=hp.PretrainedSmallLmParams(sep_token=2049),
    )

    model = model_params.instantiate().to(device)

    context = [
        torch.rand(1, 1, 200_000).to(device),
        # torch.rand(1, 1, 1234).to(device)
    ]

    # TEST TRAINING/VALIDATION STEP
    n_tries = 5
    for _ in range(n_tries):
        batch = {
            "target": torch.rand(1, 1, 320_000).to(device),
            "context": context,
            "beat": [
                Beat(beats=(torch.arange(18) * 16_000).long(),
                     downbeats=(torch.arange(0, 18, 4) * 16_000).long(),
                     seq_len=320_000)
            ],
            "description": [""],
        }
        t0 = time()
        loss = model.run_step(batch)
        torch.cuda.synchronize()
        t1 = time()
        print(f"training step in {t1 - t0} seconds")

    # print("MODEL COMPILED")
    # model = torch.compile(model, fullgraph=True,
    #                       backend="eager")  # type: ignore
    # n_tries = 5
    # for _ in range(n_tries):
    #     batch = {
    #         "target": torch.rand(2, 1, 1_000).to(device),
    #         "context": torch.rand(2, 1, 1_000).to(device) * 2,
    #         "style": torch.rand(2, 1, 1_000).to(device),
    #         "description": ["", ""],
    #     }
    #     t0 = time()
    #     loss = model.run_step(batch)
    #     t1 = time()
    #     print(f"training step in {t1 - t0} seconds")

    # TEST INFERENCE
    L.seed_everything(42)
    model.eval()

    # audio1 = load_audio(cfg.AUDIO_DIR / "42cpu.wav").to(device)
    # audio2 = load_audio(cfg.AUDIO_DIR / "42gpuT5cpu.wav").to(device)

    # PROMPT
    # prompt = audio1
    prompt = None

    # CONTEXT
    # context = torch.rand(2, 1, 320_000).to(device)
    # context = None
    # context = torch.cat((audio1, audio2), dim=0)
    # context = audio1.reshape(1, 1, -1)
    context = load_audio(cfg.EXP_DIR / "experiment_1" /
                         "context.wav").to(device).reshape(1, 1, -1)
    # context[0, ...] = 0

    # STYLE
    # style = torch.cat((audio1, audio2), dim=0)
    # style = torch.rand(2, 1, 320_000).to(device)
    style = None

    # DESCRIPTION
    description = [
        ""
        # "lo-fi chill beat with drums, keyboard and bass playing in a relaxed mood",
        # "lo-fi chill beat with drums, keyboard and bass playing in a relaxed mood"
    ]
    # description = None

    # with torch.autocast(device_type="cuda"):
    # t0 = time()
    # gen_audio = model.generate(
    #     n_samples=len(description),
    #     gen_seconds=10,
    #     prompt=prompt,
    #     context=context,
    #     style=style,
    #     description=description,
    #     prog_bar=True,
    # )
    # torch.cuda.synchronize()
    # t1 = time()
    # print(f"inference completed in {t1 - t0} seconds")

    # # for i in range(gen_audio.shape[0]):
    # save_audio(gen_audio, cfg.AUDIO_DIR / f"temp.wav")

    # args = {
    #     "n_samples": 1,
    #     "gen_seconds": 10,
    #     "prompt": prompt,
    #     "context": context,
    #     "style": style,
    #     "description": description,
    # }
    # i = iter((args,))

    # trainer = L.Trainer(precision="32", enable_progress_bar=False)

    # t0 = time()
    # gen_audio = trainer.predict(model, i)
    # torch.cuda.synchronize()
    # t1 = time()
    # print(f"predict completed in {t1 - t0} seconds")

    # save_audio(gen_audio[0], cfg.AUDIO_DIR / f"temp.wav")  # type: ignore