/home/developer6/anaconda3/envs/py310/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:28: UserWarning: torch.nn.utils.weight_norm is deprecated in favor of torch.nn.utils.parametrizations.weight_norm.
  warnings.warn("torch.nn.utils.weight_norm is deprecated in favor of torch.nn.utils.parametrizations.weight_norm.")
/home/developer6/.local/lib/python3.10/site-packages/transformers/models/encodec/modeling_encodec.py:120: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  self.register_buffer("padding_total", torch.tensor(kernel_size - stride, dtype=torch.int64), persistent=False)
Traceback (most recent call last):
  File "/home/developer6/kotoba-voice/fam/llm/train.py", line 333, in <module>
    model = KotobaSpeechModelFirstStage(config_first_stage, config_second_stage, device=training_args.device, use_kv_cache=training_args.use_kv_cache, is_debug=is_debug, configs=training_args).load_from_checkpoint(training_args.ckpt_path)
  File "/home/developer6/.local/lib/python3.10/site-packages/pytorch_lightning/utilities/model_helpers.py", line 121, in wrapper
    raise TypeError(
TypeError: The classmethod `KotobaSpeechModelFirstStage.load_from_checkpoint` cannot be called on an instance. Please call it on the class type and make sure the return value is used. import os
import threading
import json
from fam.llm import training
import torch
import wandb
import numpy as np
from scipy.io import wavfile
from torch import nn
from torch.nn import functional as F
import pytorch_lightning as pl
from pytorch_lightning import Trainer
from pytorch_lightning.strategies import FSDPStrategy, DDPStrategy
from pytorch_lightning.plugins.environments import ClusterEnvironment
from huggingface_hub import snapshot_download, HfApi, HfFolder

# Specific to your project's structure
from fam.llm.training import parse_args, get_first_stage_path, get_second_stage_path, TrainingConfig, VoiceDataModule, WandbLogger, dist_utils, optimizer_utils, Evaluator, generate_audios, upload_file_async
from fam.llm.decoders import Decoder, EncodecDecoder
from fam.quantiser.text.tokenise import TrainedBPETokeniser
from fam.llm.adapters import FlattenedInterleavedEncodec2Codebook
from fam.llm.sample import Model
import datetime
import uuid

class MyClusterEnvironment(ClusterEnvironment):
    @property
    def creates_processes_externally(self) -> bool:
        """Return True if the cluster is managed (you don't launch processes yourself)"""
        return True

    def world_size(self) -> int:
        return int(os.environ["OMPI_COMM_WORLD_SIZE"])

    def global_rank(self) -> int:
        return int(os.environ["OMPI_COMM_WORLD_RANK"])

    def local_rank(self) -> int:
        return int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"])

    def node_rank(self) -> int:
        return int(os.environ["OMPI_COMM_WORLD_NODE_RANK"])

    main_address = os.getenv("MASTER_ADDR","")

    main_port = int(os.getenv("MASTER_PORT", "0"))
    
    def set_global_rank(self, rank):
        self.global_rank_ = rank

    def set_world_size(self, size):
        self.world_size_ = size

    def detect(self):
        return True

class KotobaSpeechModelFirstStage(pl.LightningModule):
    """
    A PyTorch Lightning module for the first stage of KotobaVoice model training.
    """
    
    def __init__(self, config_first_stage, config_second_stage, device, use_kv_cache, is_debug=False, configs=None):
        super().__init__()
        self.configs = vars(configs)
        self.config_first_stage = config_first_stage
        self.use_kv_cache = use_kv_cache
        self.prev_step = -1
        self.evaluator = Evaluator()
        self.train_evaluator = Evaluator()
        self.ignore_index = -1
        self.max_seq_len = config_first_stage.block_size
        self.total_tokens = 0
        self.total_audio_tokens = 0
        self._wandb_initialized = False
        self._hf_initialized = False
        self._model_initialized = False

    def configure_model(self):
        """
        Configures the model and its components.
        """
        if self._model_initialized:
            return
        self.data_adapter = FlattenedInterleavedEncodec2Codebook(end_of_audio_token=1024)
        self.first_stage_model_cls = Model(
            self.config_first_stage,
            TrainedBPETokeniser,
            EncodecDecoder,
            data_adapter_fn=self.data_adapter.decode,
            use_kv_cache=self.use_kv_cache,
        )
        self.first_stage_model_transformer = self.first_stage_model_cls.model
        self._model_initialized = True

    def forward(self, text_tokens, embedding, inputs, targets):
        truncated_inputs = inputs[:,:self.max_seq_len] 
        truncated_first_targets = targets[:,:self.max_seq_len]
        truncated_inputs = truncated_inputs.unsqueeze(1)
        truncated_targets = truncated_first_targets.unsqueeze(1)
        # [batch_size, 7 (1 for first model, 6 for second), seq_len]

        list_logits, first_stage_unreduced_loss = self.first_stage_model_transformer(idx=truncated_inputs, embedding=embedding, targets=truncated_targets, loss_reduce="none")

        # Calculate the mean loss, excluding the ignored index (-1)
        # The operation inside the loss_reduce=mean calculation masks the -1 values 
        valid_mask = (truncated_targets != self.ignore_index)
        first_stage_loss = torch.mean(first_stage_unreduced_loss[valid_mask])

        # Calculate batch losses, excluding the ignored index for each batch
        first_stage_batch_losses = torch.tensor([
            torch.mean(first_stage_unreduced_loss[b][valid_mask[b]])
            for b in range(first_stage_unreduced_loss.size(0))
        ])

        # Asserting if the loss calculation above is the equivalent of PyTorch reduce.
        # _, assert_first_stage_loss = self.first_stage_model_transformer(idx=truncated_inputs, embedding=embedding, targets=truncated_targets, loss_reduce="mean")
        # assert(torch.isclose(first_stage_loss, assert_first_stage_loss, 1e-5))

        return first_stage_loss, first_stage_batch_losses

    def training_step(self, batch, batch_idx):
        tokens, embedding, audio_tokens, first_stage_input, first_stage_output, dataset_ids = batch 
        embedding = embedding.unsqueeze(1)

        loss, batch_losses = self(text_tokens=tokens, embedding=embedding, inputs=first_stage_input, targets=first_stage_output)

        if self.wandb_logger is not None:
            num_tokens = torch.sum(tokens != 0, 1)
            num_audio_tokens = torch.sum(first_stage_output != self.ignore_index, 1) + torch.sum(first_stage_output != self.ignore_index, 1)
            stats = {
                "dataset_id": dataset_ids,
                "num_tokens": num_tokens.tolist(),
                "num_audio_tokens": num_audio_tokens.tolist(),
                "loss": batch_losses.tolist(),
            }
            self.train_evaluator.update(stats)

        return loss
    
    def on_before_zero_grad(self, optimizer):
        if self.wandb_logger is not None:
            if self.global_step == 0:
                pass
            else:
                if dist_utils.is_main_process():
                    summary, dataset_summary = self.train_evaluator.summarize()
                    num_tokens = summary.pop("num_tokens")
                    num_audio_tokens = summary.pop("num_audio_tokens")
                    self.total_tokens += num_tokens
                    self.total_audio_tokens += num_audio_tokens

                    self.wandb_logger.log(results=summary, split="training", step=self.global_step, commit=False)
                    #self.wandb_logger.log(results=dataset_summary, split="training (details)", step=self.global_step, commit=False)
                    other_stats = {
                        f"lr_group{i}": list(self.optimizers().param_groups)[i]["lr"] for i in range(len(self.optimizers().param_groups))
                    }
                    other_stats["tokens"] = self.total_tokens
                    other_stats["audio_tokens"] = self.total_audio_tokens
                    other_stats["examples"] = self.configs["batch_size"] * self.global_step
                    self.wandb_logger.log(results=other_stats, split="stats", step=self.global_step, commit=False)

                self.train_evaluator.reset()


    def on_train_start(self): 
        if dist_utils.is_main_process():
            if self.configs["use_wandb"] and not self._wandb_initialized:
                self.wandb_logger = WandbLogger(training_args)
                wandb.watch(self.first_stage_model_transformer, log='parameters', log_freq=1000)
                self._wandb_initialized = True
            if self.configs["upload_checkpoint_to_hf"] and not self._hf_initialized:
                wandb_run_id = self.wandb_logger.get_run_id()
                self._hf_repo_name = "kotoba-speech/{}_{}".format(self.configs["exp_name"], wandb_run_id)
                api = HfApi()
                hf_token = os.environ.get('HF_TOKEN')
                api.create_repo(token=hf_token, repo_id=self._hf_repo_name, repo_type="model", exist_ok=True, private=True)
                self._hf_initialized = True
        else:
            self.wandb_logger = None
    
    def on_validation_start(self): 
        if dist_utils.is_main_process():
            if self.configs["use_wandb"] and not self._wandb_initialized:
                self.wandb_logger = WandbLogger(training_args)
                wandb.watch(self.first_stage_model_transformer, log='parameters', log_freq=1000)
                self._wandb_initialized = True
            if self.configs["upload_checkpoint_to_hf"] and not self._hf_initialized:
                wandb_run_id = self.wandb_logger.get_run_id()
                self._hf_repo_name = "kotoba-speech/{}_{}".format(self.configs["exp_name"], wandb_run_id)
                api = HfApi()
                hf_token = os.environ.get('HF_TOKEN')
                api.create_repo(token=hf_token, repo_id=self._hf_repo_name, repo_type="model", exist_ok=True, private=True)
                self._hf_initialized = True
        else:
            self.wandb_logger = None
            
    def validation_step(self, batch, batch_idx):
        tokens, embedding, audio_tokens, first_stage_input, first_stage_output, dataset_ids = batch
        tokens = tokens.unsqueeze(1)
        embedding = embedding.unsqueeze(1)
        loss, batch_losses = self(text_tokens=tokens, embedding=embedding, inputs=first_stage_input, targets=first_stage_output)

        stats = {
            "dataset_id": dataset_ids,
            "loss": batch_losses.tolist(),
        }
        self.evaluator.update(stats)
        return loss

    def on_validation_epoch_end(self):
        self.evaluator.synchronize_between_processes()
        if dist_utils.is_main_process():
            summary, dataset_summary = self.evaluator.summarize()
        self.evaluator.reset()
        if dist_utils.is_main_process():
            if hasattr(self, 'wandb_logger'):
                self.wandb_logger.log(results=summary, split="validation", step=self.global_step, commit=False)
                self.wandb_logger.log(results=dataset_summary, split="validation (details)", step=self.global_step, commit=False)    

    def on_validation_end(self):
        if dist_utils.is_main_process():
            if self.configs["upload_checkpoint_to_hf"]:
                if self.trainer.checkpoint_callback.last_model_path != "":
                    chpt_dir_name = self.trainer.checkpoint_callback.state_dict()["dirpath"]
                    files = [os.path.join(chpt_dir_name, f) for f in os.listdir(chpt_dir_name) if f.endswith(".ckpt") and f != "last.ckpt"]
                    newest_file = max(files, key=os.path.getmtime)
                    try:
                        upload_thread = threading.Thread(target=upload_file_async, args=(newest_file, os.path.basename(newest_file), self._hf_repo_name, "model"))
                        upload_thread.start()
                    except:
                        print("HF upload failed for checkpoint: {}".format(newest_file))

    def on_save_checkpoint(self, checkpoint):
        if dist_utils.is_main_process():
            # TODO: Fix this. This is probably saving checkpoints twice. Ideally, overwriting on_save_checkpoint in the callback is cleaner.
            if self.configs["log_audio_sample"]:
                with torch.no_grad():
                    ckpt_path = os.path.join("lightning_logs", uuid.uuid4().hex + ".pt")
                    torch.save(checkpoint, ckpt_path)
                    print("Generating sample audios for checkpoint : {}".format(ckpt_path))
                    results = generate_audios(first_model_path=ckpt_path, file_path=self.configs["validation_file_path"], guidance_scale=self.configs["guidance_scale"])
                    audio_files = {}
                    for i, rs in enumerate(results):
                        id = rs["id"] if "id" in rs else "ex-{}".format(i)
                        caption = rs["prompt"]
                        wavpath = rs["audio_path"]
                        sample_rate, data = wavfile.read(wavpath)
                        if data.dtype != np.float32:
                            data = data.astype(np.float32) / np.max(np.abs(data))  # Normalize data to [-1, 1]
                        audio_files[id] = wandb.Audio(data, caption=caption, sample_rate=sample_rate)
                    self.wandb_logger.log(results=audio_files, split="audio", step=self.global_step, commit=False)
                    os.remove(ckpt_path)

    def configure_optimizers(self):
        return optimizer_utils.set_schedule(self)

# Train the model
if __name__ == "__main__":
    training_args = parse_args()
    data_module = VoiceDataModule(training_args.per_gpu_batchsize, training_args.data_list, training_args.local_path, training_args.ratio_filter, training_args.per_video_spk)

    # assertions
    if training_args.upload_checkpoint_to_hf:
        assert training_args.use_wandb, "HF uploading requires wandb run_id, so use_wandb must be enabled."
        # Validate HF token
        api = HfApi()
        try:
            hf_token = os.environ.get('HF_TOKEN')
            if not hf_token:
                raise ValueError("HF_TOKEN environment variable is not set or empty. Please export HF_TOKEN by running 'export HF_TOKEN=YOUR_TOKEN'.")

            user_info = api.whoami(token=hf_token)
            print("User info retrieved:", user_info)
        except Exception as e:
            raise ValueError("Invalid HF token provided or error retrieving user info. Please ensure the HF_TOKEN is correctly set.") from e

    print('ckpt path in main', training_args.ckpt_path)

    model_dir = snapshot_download(repo_id=training_args.huggingface_repo_id)
    first_stage_ckpt_path = get_first_stage_path(model_dir)
    second_stage_ckpt_path = get_second_stage_path(model_dir)
    config_first_stage = TrainingConfig(
        ckpt_path=first_stage_ckpt_path,
        num_samples=training_args.num_samples,
        seed=training_args.seed,
        device=training_args.device,
        dtype=training_args.dtype,
        compile=training_args.compile,
        init_from=training_args.init_from,
        train_from_scratch=training_args.train_from_scratch,
        output_dir=training_args.output_dir,
        spkemb_dropout=training_args.spkemb_dropout,
        block_size=training_args.block_size,
        block_size_override=training_args.block_size_override #linear interpolates existing positional embeddings to extend CL.
    )

    config_second_stage = TrainingConfig(
        ckpt_path=second_stage_ckpt_path,
        num_samples=training_args.num_samples,
        seed=training_args.seed,
        device=training_args.device,
        dtype=training_args.dtype,
        compile=training_args.compile,
        init_from=training_args.init_from,
        train_from_scratch=training_args.train_from_scratch,
        output_dir=training_args.output_dir,
        spkemb_dropout=training_args.spkemb_dropout,
        block_size=training_args.block_size,
        block_size_override=training_args.block_size_override
    )

    is_debug = training_args.debug
    model = KotobaSpeechModelFirstStage(config_first_stage, config_second_stage, device=training_args.device, use_kv_cache=training_args.use_kv_cache, is_debug=is_debug, configs=training_args)
    #model.predict_step()

    # # #

    #checkpoint and finetuning loading

    # # #

    if training_args.finetuning_ckpt_path is not None:
        ckpt = torch.load(training_args.finetuning_ckpt_path)
        state_dict = ckpt["state_dict"]
        model.configure_model()
        model.load_state_dict(state_dict)
    lr_callback = pl.callbacks.LearningRateMonitor(logging_interval="step")

    if training_args.ckpt_path is not None:
        ckpt = torch.load(training_args.ckpt_path)
        state_dict = ckpt["state_dict"]
        model.configure_model()    
        model = KotobaSpeechModelFirstStage(config_first_stage, config_second_stage, device=training_args.device, use_kv_cache=training_args.use_kv_cache, is_debug=is_debug, configs=training_args).load_from_checkpoint(training_args.ckpt_path)