training/callback.py

from concurrent.futures import ThreadPoolExecutor
import pickle
import librosa
import lightning as L
from lightning.pytorch.utilities.seed import isolate_rng
from lightning.pytorch.callbacks import Callback
from torch import Tensor
from typing import Dict, Any
from pathlib import Path
import random
import wandb
import torch
from lightning.pytorch.utilities import rank_zero_only

from conditioning.beat_embedder import Beat
from utils import audio
import config as cfg
# from cocola.contrastive_model import CoCola
# from cocola.feature_extraction import CoColaFeatureExtractor
# from cocola import constants as cocola_constants
# from utils.logging import upload_to_s3

# Custom progress bar to refresh only twice per epoch
# class SlowProgressBar(TQDMProgressBar):

#     def init_train_tqdm(self):
#         bar = super().init_train_tqdm()
#         # Refresh rate: update the bar only twice per epoch
#         bar.refresh_rate = self.total_train_batches // 2
#         return bar


class SaveDemoOnValidationCallback(Callback):

    @rank_zero_only
    def __init__(self, save_path: Path, save_model: bool, n_demos: int):
        self.save_path: Path = save_path
        self.save_model: bool = save_model
        self.save_path.mkdir(parents=True, exist_ok=True)
        self.n_demos: int = n_demos
        self.executor = ThreadPoolExecutor(max_workers=3)

    # @rank_zero_only
    # def on_validation_end(self, trainer: L.Trainer,
    #                       pl_module: L.LightningModule) -> None:
    #     if self.save_model:
    #         if not (self.save_path / "params.pkl").exists():
    #             with open(self.save_path / "params.pkl", "wb") as f:
    #                 pickle.dump(pl_module.params, f)

    #         checkpoint_path = self.save_path / f"step={trainer.global_step}.ckpt"
    #         trainer.save_checkpoint(checkpoint_path, weights_only=False)

    #         # upload checkpoint to S3
    #         run_name = self.save_path.name
    #         future = self.executor.submit(
    #             upload_to_s3,
    #             checkpoint_path,
    #             "lag-modular",
    #             f"checkpoints/{run_name}/{checkpoint_path.name}",
    #         )

    #         def log_upload_status(fut):
    #             try:
    #                 success = fut.result()
    #                 status = "successful" if success else "failed"
    #                 print(f"Checkpoint upload on S3 {status}")
    #             except Exception as e:
    #                 print(f"Error during upload of {checkpoint_path}: {e}")

    #         future.add_done_callback(log_upload_status)

    @rank_zero_only
    def on_validation_epoch_start(self, trainer: L.Trainer,
                                  pl_module: L.LightningModule) -> None:

        total_batches = trainer.num_val_batches[0]
        assert isinstance(total_batches, int)
        self.random_batch_idx = random.randint(0, total_batches - 1)

    def on_train_end(self, trainer, pl_module):
        print("Shutting down uploader")
        self.executor.shutdown(wait=True)
        print("Uploader shut down")

    @rank_zero_only
    def on_validation_batch_end(self,
                                trainer: L.Trainer,
                                pl_module: L.LightningModule,
                                outputs: Any,
                                batch: Dict[str, Any],
                                batch_idx: int,
                                dataloader_idx: int = 0) -> None:
        if batch_idx == self.random_batch_idx:
            self.generate_and_save_demo(trainer, pl_module, batch)

    def generate_and_save_demo(
        self,
        trainer: L.Trainer,
        pl_module: L.LightningModule,
        batch: Dict[str, Any],
    ):
        logger = trainer.logger.experiment if trainer.logger is not None else None  # type: ignore

        save_path: Path = self.save_path / "demos" / f"step={trainer.global_step}"
        save_path.mkdir(parents=True, exist_ok=True)

        # batch_size = min(len(batch["target"]), self.n_demos)
        batch_size = len(batch["target"])
        if self.n_demos < len(batch["target"]):
            chosen_samples = random.sample(list(range(batch_size)),
                                           self.n_demos)
        else:
            chosen_samples = list(range(batch_size))
        n_samples = len(chosen_samples)

        with isolate_rng(include_cuda=True):
            # extract conditioning data from batch
            target = batch.get("target")
            assert target is not None
            context = batch.get("context")
            context_dropout_mask = None
            if context is not None:
                if isinstance(context, list):
                    context = [
                        c for i, c in enumerate(context) if i in chosen_samples
                    ]
                else:
                    assert isinstance(context, Tensor)
                    context = context[chosen_samples]
                context_dropout_mask = torch.full((len(context),),
                                                  True,
                                                  dtype=torch.bool)
                # context_dropout_mask[:n_samples // 2] = False # todo uncomment me
            style = batch.get("style")
            if style is not None:
                style = style[chosen_samples]
            description = batch.get("description")
            if description is not None:
                description = [
                    v for i, v in enumerate(description) if i in chosen_samples
                ]
            beat = batch.get("beat")
            if beat is not None:
                beat = [v for i, v in enumerate(beat) if i in chosen_samples]

            # generate
            L.seed_everything(42)
            with torch.autocast(device_type="cuda"):
                gen_audio = pl_module.generate(
                    n_samples=n_samples,
                    gen_seconds=10,
                    prompt=None,
                    context=context,
                    context_dropout_mask=context_dropout_mask,
                    style=style,
                    beat=beat,
                    description=description,
                    prog_bar=cfg.running_locally())

            # compute cocola score between generated audio and context/style
            '''if style is not None or context is not None:
                cocola = CoCola(
                    embedding_mode=cocola_constants.EmbeddingMode.BOTH)
                cocola.load_state_dict(
                    torch.load(cfg.weights_dir() / "cocola-weights.pt",
                               weights_only=True))
                # cocola = CoCola.load_from_checkpoint(cfg.weights_dir() /
                #                                      "cocola.pt",
                #                                      map_location="cpu").eval()
                cocola.eval()
                cocola.set_embedding_mode(cocola_constants.EmbeddingMode.BOTH)
                feature_extractor = CoColaFeatureExtractor()
                gen_features = feature_extractor(gen_audio.cpu())
                if style is not None:
                    assert len(gen_audio) == len(style)
                    style_features = feature_extractor(style.cpu())
                    style_score = cocola.score(gen_features, style_features)
                if context is not None and isinstance(context, Tensor):
                    assert len(gen_audio) == len(context)
                    context_features = feature_extractor(context.cpu())
                    context_score = cocola.score(gen_features, context_features)'''

        columns = ([
            s for s in ("description", "context", "beat") if s in batch
        ] + ["generated"])
        if "context" in columns:
            columns.append("mix")
            # mixed_audio = gen_audio + context
        if "beat" in columns and "context" in columns:
            columns.append("context with beat")
        # if "style" in columns:
        #     columns.append("style cocola score")
        # if "context" in columns and isinstance(context, Tensor):
        #     columns.append("context cocola score")

        data = [[] for _ in range(n_samples)]

        # for idx in range(n_samples):
        for i, idx in enumerate(chosen_samples):
            gen_filename: Path = save_path / f"demo{i}_gen.wav"
            audio.save_audio(gen_audio[i], gen_filename)

            if description is not None:
                desc_filename: Path = save_path / f"demo{i}_description.txt"
                desc_filename.write_text(description[i])
            '''if style is not None:
                style_filename: Path = save_path / f"demo{i}_style.wav"
                audio.save_audio(style[i], style_filename)'''
            if context is not None and context_dropout_mask[i]:
                context_filename: Path = save_path / f"demo{i}_context.wav"
                mix_filename: Path = save_path / f"demo{i}_mix.wav"
                audio.save_audio(context[i], context_filename)
                mix = torch.nn.functional.pad(
                    context[i],
                    (0, gen_audio[i].shape[-1] - context[i].shape[-1]),
                    value=0) + gen_audio[i]
                audio.save_audio(mix, mix_filename)
            if "target" in batch:
                target_filename: Path = save_path / f"demo{i}_target.wav"
                audio.save_audio(target[idx], target_filename)

            if beat is not None:
                beat_filename: Path = save_path / f"demo{i}_beat.wav"
                b: Beat = beat[i]
                beats = (b.beats / pl_module.sample_rate).cpu().numpy()
                downbeats = (b.downbeats / pl_module.sample_rate).cpu().numpy()
                clicks = librosa.clicks(times=beats,
                                        sr=pl_module.sample_rate,
                                        click_freq=1000,
                                        length=gen_audio[i].shape[-1])
                high_clicks = librosa.clicks(times=downbeats,
                                             sr=pl_module.sample_rate,
                                             click_freq=2000,
                                             length=gen_audio[i].shape[-1])
                gen_with_beat = gen_audio[i].cpu() + (clicks *
                                                      0.5) + (high_clicks * 0.5)
                audio.save_audio(gen_with_beat, beat_filename)
                if context is not None and context_dropout_mask[i]:
                    context_with_beats_filename: Path = (
                        save_path / f"demo{i}_context_with_beats.wav")
                    context_with_beats = torch.nn.functional.pad(
                        context[i].cpu(),
                        (0, len(clicks) - context[i].shape[-1])) + (
                            clicks * 0.5) + (high_clicks * 0.5)
                    audio.save_audio(context_with_beats,
                                     context_with_beats_filename)

            if logger is not None:
                # if pl_module.logger is not None:
                row = []
                if description is not None:
                    row.append(description[i])
                '''if style is not None:
                    row.append(
                        wandb.Audio(str(style_filename),
                                    sample_rate=pl_module.sample_rate))'''
                if context is not None:
                    assert context_dropout_mask is not None
                    if context_dropout_mask[i]:
                        row.append(
                            wandb.Audio(str(context_filename),
                                        sample_rate=pl_module.sample_rate))
                    else:
                        row.append(None)

                if beat is not None:
                    row.append(
                        wandb.Audio(str(beat_filename),
                                    sample_rate=pl_module.sample_rate))

                row.append(
                    wandb.Audio(str(gen_filename),
                                sample_rate=pl_module.sample_rate))
                if context is not None:
                    assert context_dropout_mask is not None
                    if context_dropout_mask[i]:
                        row.append(
                            wandb.Audio(str(mix_filename),
                                        sample_rate=pl_module.sample_rate))
                    else:
                        row.append(None)
                if context is not None and beat is not None:
                    assert context_dropout_mask is not None
                    if context_dropout_mask[i]:
                        row.append(
                            wandb.Audio(str(context_with_beats_filename),
                                        sample_rate=pl_module.sample_rate))
                    else:
                        row.append(None)
                # if style is not None:
                #     row.append(style_score[i])
                # if context is not None and isinstance(context, Tensor):
                #     row.append(context_score[i])

                data[i] = row

        # if pl_module.logger is not None:
        if logger is not None:
            table = wandb.Table(columns=columns, data=data)
            logger.log({f"demos/step:{trainer.global_step}": table})
            # if style is not None:
            #     logger.log({"valid/style_cocola": style_score.mean()})
            # if context is not None and isinstance(context, Tensor):
            #     logger.log({"valid/context_cocola": context_score.mean()})
            # logger.log_table(  # type: ignore
            #     key=f"demos_step{trainer.global_step}",
            #     columns=columns,
            #     data=data)


class PrintLossesCallback(L.Callback):

    def __init__(self):
        super().__init__()
        self.train_losses = []
        self.val_losses = []

    def on_train_epoch_end(self, trainer, pl_module):
        # Compute and print the average training loss for the epoch
        train_loss = trainer.callback_metrics.get("train_loss")
        if train_loss is None:
            return
        print(f"Epoch {trainer.current_epoch +1} / "
              f"train loss: {train_loss:.4f}")

    def on_validation_epoch_end(self, trainer, pl_module):
        # Compute and print the average validation loss for the epoch
        val_loss = trainer.callback_metrics.get("val_loss")
        if val_loss is None:
            return
        print(f"Epoch {trainer.current_epoch + 1} / "
              f"validation loss: {val_loss:.4f}")