"""PyTorch Lightning module for InstructMusicGen model with adapter-based instruction tuning.

Trimmed to the inference path only i.e. removed training_step/ validation_step/test_step, torchmetrics/wandb 
logging bcs not needed for inference; only kept what load_from_checkpoint() needs.
"""
import torch
from lightning import LightningModule
from .components.model import Instructor


class InstructMusicGenAdapterLitModule(LightningModule):
    def __init__(
            self,
            optimizer: torch.optim.Optimizer,
            scheduler: torch.optim.lr_scheduler,
            tmp_dir: str,
            compile: bool,
            instructor: Instructor,
            audio_regularization: float,
    ) -> None:

        super().__init__()

        # this line allows to access init params with 'self.hparams' attribute
        # also ensures init params will be stored in ckpt
        self.save_hyperparameters(logger=False)

        self.model = self.hparams.instructor()