docstr: Tortoise GPT training class
docstr_tok: ["Tortoise","GPT","training","class"]
code: def __init__(self, config: Coqpit):
      """
      Tortoise GPT training class
      """
      super().__init__(config, ap=None, tokenizer=None)
      self.config = config
      # init XTTS model
      self.xtts = Xtts(self.config)
      # create the tokenizer with the target vocabulary
      self.xtts.tokenizer = VoiceBpeTokenizer(self.args.tokenizer_file)
      # init gpt encoder and hifigan decoder
      self.xtts.init_models()
  
      if self.args.xtts_checkpoint:
          self.load_checkpoint(
              self.config,
              self.args.xtts_checkpoint,
              eval=False,
              strict=False,
          )
  
      # set mel stats
      if self.args.mel_norm_file:
          self.xtts.mel_stats = load_fsspec(self.args.mel_norm_file)
  
      # load GPT if available
      if self.args.gpt_checkpoint:
          gpt_checkpoint = torch.load(
              self.args.gpt_checkpoint, map_location=torch.device("cpu")
          )
          # deal with coqui Trainer exported model
          if (
              "model" in gpt_checkpoint.keys()
              and "config" in gpt_checkpoint.keys()
          ):
              print("Coqui Trainer checkpoint detected! Converting it!")
              gpt_checkpoint = gpt_checkpoint["model"]
              states_keys = list(gpt_checkpoint.keys())
              for key in states_keys:
                  if "gpt." in key:
                      new_key = key.replace("gpt.", "")
                      gpt_checkpoint[new_key] = gpt_checkpoint[key]
                      del gpt_checkpoint[key]
                  else:
                      del gpt_checkpoint[key]
  
          # edit checkpoint if the number of tokens is changed to ensures the better transfer learning possible
          if (
              "text_embedding.weight" in gpt_checkpoint
              and gpt_checkpoint["text_embedding.weight"].shape
              != self.xtts.gpt.text_embedding.weight.shape
          ):
              num_new_tokens = (
                  self.xtts.gpt.text_embedding.weight.shape[0]
                  - gpt_checkpoint["text_embedding.weight"].shape[0]
              )
              print(
                  f" > Loading checkpoint with {num_new_tokens} additional tokens."
              )
  
              # add new tokens to a linear layer (text_head)
              emb_g = gpt_checkpoint["text_embedding.weight"]
              new_row = torch.randn(num_new_tokens, emb_g.shape[1])
              start_token_row = emb_g[-1, :]
              emb_g = torch.cat([emb_g, new_row], axis=0)
              emb_g[-1, :] = start_token_row
              gpt_checkpoint["text_embedding.weight"] = emb_g
  
              # add new weights to the linear layer (text_head)
              text_head_weight = gpt_checkpoint["text_head.weight"]
              start_token_row = text_head_weight[-1, :]
              new_entry = torch.randn(
                  num_new_tokens,
                  self.xtts.gpt.text_head.weight.shape[1],
              )
              text_head_weight = torch.cat(
                  [text_head_weight, new_entry], axis=0
              )
              text_head_weight[-1, :] = start_token_row
              gpt_checkpoint["text_head.weight"] = text_head_weight
  
              # add new biases to the linear layer (text_head)
              text_head_bias = gpt_checkpoint["text_head.bias"]
              start_token_row = text_head_bias[-1]
              new_bias_entry = torch.zeros(num_new_tokens)
              text_head_bias = torch.cat(
                  [text_head_bias, new_bias_entry], axis=0
              )
              text_head_bias[-1] = start_token_row
              gpt_checkpoint["text_head.bias"] = text_head_bias
  
          self.xtts.gpt.load_state_dict(gpt_checkpoint, strict=True)
          print(
              ">> GPT weights restored from:", self.args.gpt_checkpoint
          )
  
      # Mel spectrogram extractor for conditioning
      if self.args.gpt_use_perceiver_resampler:
          self.torch_mel_spectrogram_style_encoder = (
              TorchMelSpectrogram(
                  filter_length=2048,
                  hop_length=256,
                  win_length=1024,
                  normalize=False,
                  sampling_rate=config.audio.sample_rate,
                  mel_fmin=0,
                  mel_fmax=8000,
                  n_mel_channels=80,
                  mel_norm_file=self.args.mel_norm_file,
              )
          )
      else:
          self.torch_mel_spectrogram_style_encoder = (
              TorchMelSpectrogram(
                  filter_length=4096,
                  hop_length=1024,
                  win_length=4096,
                  normalize=False,
                  sampling_rate=config.audio.sample_rate,
                  mel_fmin=0,
                  mel_fmax=8000,
                  n_mel_channels=80,
                  mel_norm_file=self.args.mel_norm_file,
              )
          )
  
      # Load DVAE
      self.dvae = DiscreteVAE(
          channels=80,
          normalization=None,
          positional_dims=1,
          num_tokens=self.args.gpt_num_audio_tokens - 2,
          codebook_dim=512,
          hidden_dim=512,
          num_resnet_blocks=3,
          kernel_size=3,
          num_layers=2,
          use_transposed_convs=False,
      )
  
      self.dvae.eval()
      if self.args.dvae_checkpoint:
          dvae_checkpoint = torch.load(
              self.args.dvae_checkpoint,
              map_location=torch.device("cpu"),
          )
          self.dvae.load_state_dict(dvae_checkpoint, strict=False)
          print(
              ">> DVAE weights restored from:",
              self.args.dvae_checkpoint,
          )
      else:
          raise RuntimeError(
              "You need to specify config.model_args.dvae_checkpoint path to be able to train the GPT decoder!!"
          )
  
      # Mel spectrogram extractor for DVAE
      self.torch_mel_spectrogram_dvae = TorchMelSpectrogram(
          mel_norm_file=self.args.mel_norm_file,
          sampling_rate=config.audio.dvae_sample_rate,
      )
  
imports: from dataclasses import dataclass, field
  from typing import Dict, List, Tuple, Union
  import torch
  import torch.nn
  import torchaudio
  from coqpit import Coqpit
  from torch.nn import functional
  from torch.utils.data import DataLoader
  from trainer.torch import DistributedSampler
  from trainer.trainer_utils import get_optimizer, get_scheduler
  from TTS.tts.configs.xtts_config import XttsConfig
  from TTS.tts.datasets.dataset import TTSDataset
  from TTS.tts.layers.tortoise.arch_utils import TorchMelSpectrogram
  from TTS.tts.layers.xtts.dvae import DiscreteVAE
  from TTS.tts.layers.xtts.tokenizer import VoiceBpeTokenizer
  from TTS.tts.layers.xtts.trainer.dataset import XTTSDataset
  from TTS.tts.models.base_tts import BaseTTS
  from TTS.tts.models.xtts import Xtts, XttsArgs, XttsAudioConfig
  from TTS.utils.io import load_fsspec
label_desc: Use image node_id 1 to create a new GPTTrainer object from inherited base classes: BaseTTS with example: obj = GPTTrainer(config)