content: from __future__ import annotations
  
  import math
  from functools import partial, wraps
  
  from beartype import beartype
  
  import torch
  from torch import nn, einsum, Tensor
  from torch.autograd import grad as torch_grad
  import torch.nn.functional as F
  from torch.nn.utils.rnn import pad_sequence
  
  import torchaudio
  
  from einops import rearrange, repeat, reduce
  from einops.layers.torch import Rearrange
  
  from audiolm_pytorch.vq_wav2vec import FairseqVQWav2Vec
  from audiolm_pytorch.hubert_kmeans import HubertWithKmeans
  
  from audiolm_pytorch.t5 import t5_encode_text, get_encoded_dim, DEFAULT_T5_NAME
  
  from hyper_connections import get_init_and_expand_reduce_stream_functions
  
  from torchaudio.functional import resample
  
  from audiolm_pytorch.soundstream import SoundStream
  from audiolm_pytorch.encodec import EncodecWrapper
  from audiolm_pytorch.utils import AudioConditionerBase
  from audiolm_pytorch.attend import Attend
  
  from tqdm import tqdm
  from pathlib import Path
  from audiolm_pytorch.version import __version__
  from packaging import version
  
  # helper functions
  
  def exists(val):
      return val is not None
  
  def default(val, d):
      return val if exists(val) else d
  
  def always(val):
      def inner(*args, **kwargs):
          return val
      return inner
  
  def maybe(fn):
      if not exists(fn):
          return always(None)
  
      @wraps(fn)
      def inner(x, *args, **kwargs):
          if not exists(x):
              return x
          return fn(x, *args, **kwargs)
      return inner
  
  def ceil_div(numer, denom):
      return (numer + denom - 1) // denom
  
  def remainder_needed_until_multiple(n, mult):
      return (ceil_div(n, mult) * mult) - n
  
  def round_down_nearest_multiple(val, mult):
      return (val // mult) * mult
  
  def eval_decorator(fn):
      def inner(model, *args, **kwargs):
          was_training = model.training
          model.eval()
          out = fn(model, *args, **kwargs)
          model.train(was_training)
          return out
      return inner
  
  # tensor helpers
  
  def generate_mask_with_prob(shape, mask_prob, device):
      seq = shape[-1]
      rand = torch.randn(shape, device = device)
      rand[:, 0] = -torch.finfo(rand.dtype).max
      num_mask = min(int(seq * mask_prob), seq - 1)
      indices = rand.topk(num_mask, dim = -1).indices
      mask = ~torch.zeros(shape, device = device).scatter(1, indices, 1.).bool()
      return mask
  
  # attention related utils
  
  def grad_shrink(t, alpha = 0.1):
      return t * alpha + t.detach() * (1 - alpha)
  
  # sampling helpers
  
  def log(t, eps = 1e-20):
      return torch.log(t + eps)
  
  def l2norm(t):
      return F.normalize(t, dim = -1)
  
  def gumbel_noise(t):
      noise = torch.zeros_like(t).uniform_(0, 1)
      return -log(-log(noise))
  
  def gumbel_sample(t, temperature = 1., dim = -1):
      return ((t / temperature) + gumbel_noise(t)).argmax(dim = dim)
  
  def top_k(logits, thres = 0.5):
      num_logits = logits.shape[-1]
      k = max(int((1 - thres) * num_logits), 1)
      val, ind = torch.topk(logits, k)
      probs = torch.full_like(logits, float('-inf'))
      probs.scatter_(1, ind, val)
      return probs
  
  def mask_out_after_eos_id(t, eos_id, mask_value = -1, keep_eos = True):
      eos_mask = (t == eos_id).float()
  
      if keep_eos:
          eos_mask = F.pad(eos_mask, (1, -1))
  
      after_eos_mask = eos_mask.cumsum(dim = -1) > 0
      return t.masked_fill(after_eos_mask, mask_value)
  
  def all_rows_have_eos_id(t, eos_id):
      eos_mask = (t == eos_id)
      return torch.any(eos_mask, dim = -1).all()
  
  def safe_cat(*tensors, dim = -2):
      args = [*filter(exists, tensors)]
  
      if len(args) == 0:
          return None
      elif len(args) == 1:
          return args[0]
      else:
          return torch.cat(args, dim = dim)
  
  # classifier free guidance functions
  
  def prob_mask_like(shape, prob, device):
      if prob == 1:
          return torch.ones(shape, device = device, dtype = torch.bool)
      elif prob == 0:
          return torch.zeros(shape, device = device, dtype = torch.bool)
      else:
          return torch.zeros(shape, device = device).float().uniform_(0, 1) < prob
  
  # removing unique consecutives in the semantic token ids
  # important detail noted by @eonglints
  
  def append_eos_id(ids, eos_id):
      b, device = ids.shape[0], ids.device
      eos_ids = torch.ones(1, device = device).long() * eos_id
      eos_ids = repeat(eos_ids, '1 -> b 1', b = b)
      ids = torch.cat((ids, eos_ids), dim = -1)
      return ids
  
  def batch_unique_consecutive(t, pad_value = 0.):
      unique_arr = [torch.unique_consecutive(el) for el in t.unbind(dim = 0)]
      return pad_sequence(unique_arr, batch_first = True, padding_value = pad_value)
  
  # function for getting embeds from nn.Embedding but with padding as some designated value (-1) outside the range of the embed table
  
  @beartype
  def get_embeds(
      embeddings: nn.Embedding,
      codes: torch.Tensor,
      pad_id = -1,
      return_mask = False,
      mask_pad_pos_to = 0
  ):
      pad_mask = codes == pad_id
      codes_without_pad = codes.masked_fill(pad_mask, 0) # just retrieve first code as dummy
      embeds = embeddings(codes_without_pad)
  
      if exists(mask_pad_pos_to):
          embeds = embeds.masked_fill(rearrange(pad_mask, '... -> ... 1'), mask_pad_pos_to)
  
      if return_mask:
          return embeds, ~pad_mask
  
      return embeds
  
  # bias-less layernorm, being used in more recent T5s, PaLM, also in @borisdayma 's experiments shared with me
  # greater stability
  
  class LayerNorm(nn.Module):
      def __init__(self, dim):
          super().__init__()
          self.gamma = nn.Parameter(torch.ones(dim))
          self.register_buffer("beta", torch.zeros(dim))
  
      def forward(self, x):
          return F.layer_norm(x, x.shape[-1:], self.gamma, self.beta)
  
  # relative positional bias
  
  class RelativePositionBias(nn.Module):
      """ from https://arxiv.org/abs/2111.09883 """
  
      def __init__(
          self,
          *,
          dim,
          heads,
          layers = 3
      ):
          super().__init__()
          self.net = nn.ModuleList([])
          self.net.append(nn.Sequential(nn.Linear(1, dim), nn.SiLU()))
  
          for _ in range(layers - 1):
              self.net.append(nn.Sequential(nn.Linear(dim, dim), nn.SiLU()))
  
          self.net.append(nn.Linear(dim, heads))
  
      @property
      def device(self):
          return next(self.parameters()).device
  
      def forward(self, i, j):
          assert j >= i
          device = self.device
  
          i_pos = torch.arange(i, device = device) + (j - i)
          j_pos = torch.arange(j, device = device)
  
          rel_pos = (rearrange(i_pos, 'i -> i 1') - rearrange(j_pos, 'j -> 1 j'))
          rel_pos += (j - 1)
  
          x = torch.arange(-j + 1, j, device = device).float()
          x = rearrange(x, '... -> ... 1')
  
          for layer in self.net:
              x = layer(x)
  
          x = x[rel_pos]
          return rearrange(x, 'i j h -> h i j')
  
  # feedforward
  
  class GEGLU(nn.Module):
      def forward(self, x):
          x, gate = x.chunk(2, dim = -1)
          return F.gelu(gate) * x
  
  def FeedForward(dim, mult = 4, dropout = 0.1):
      inner_dim = int(dim * 2 * mult / 3)
      return nn.Sequential(
          LayerNorm(dim),
          nn.Linear(dim, inner_dim * 2, bias = False),
          GEGLU(),
          LayerNorm(inner_dim),
          nn.Dropout(dropout),
          nn.Linear(inner_dim, dim, bias = False)
      )
  
  # attention
  
  class Attention(nn.Module):
      def __init__(
          self,
          dim,
          causal = False,
          dim_head = 64,
          dim_context = None,
          heads = 8,
          norm_context = False,
          num_null_kv = 0,
          dropout = 0.1,
          scale = 8,
          flash = False
      ):
          super().__init__()
          self.heads = heads
          self.causal = causal
          inner_dim = dim_head * heads
  
          dim_context = default(dim_context, dim)
  
          self.norm = LayerNorm(dim)
          self.context_norm = LayerNorm(dim_context) if norm_context else nn.Identity()
  
          self.attn_dropout = nn.Dropout(dropout)
  
          self.num_null_kv = num_null_kv
          self.null_kv = nn.Parameter(torch.randn(2, num_null_kv, dim_head)) if num_null_kv > 0 else None
  
          self.to_q = nn.Linear(dim, inner_dim, bias = False)
          self.to_kv = nn.Linear(dim_context, dim_head * 2, bias = False)
  
          self.attend = Attend(
              flash = flash,
              dropout = dropout,
              causal = causal
          )
  
          self.to_out = nn.Sequential(
              nn.Linear(inner_dim, dim, bias = False),
              nn.Dropout(dropout)
          )
  
      def forward(
          self,
          x,
          context = None,
          mask = None,
          attn_bias = None,
          prefix_context = None,
          prefix_context_mask = None,
          return_kv_cache = False,
          return_values = False,
          value_residual: Tensor | None = None,
          kv_cache = None
      ):
          b, n, _, device = *x.shape, x.device
  
          if exists(context):
              context = self.context_norm(context)
  
          kv_input = default(context, x)
  
          # take care of prefix-based self attention conditioning
          # make sure to either concat the to the self attention mask or lengthen it accordingly
  
          if exists(prefix_context):
              kv_input = torch.cat((prefix_context, kv_input), dim = -2)
              prefix_seq_len = prefix_context.shape[-2]
  
              if not exists(mask):
                  mask = torch.ones((b, n), device = device, dtype = torch.bool)
  
              if exists(prefix_context_mask):
                  mask = torch.cat((prefix_context_mask, mask), dim = -1)
              else:
                  mask = F.pad(mask, (prefix_seq_len, 0), value = True)
  
              if exists(attn_bias):
                  attn_bias = F.pad(attn_bias, (prefix_seq_len, 0), value = 0.)
  
          # prenorm
  
          x = self.norm(x)
  
          # project for queries, keys, values
  
          q, k, v = self.to_q(x), *self.to_kv(kv_input).chunk(2, dim = -1)
  
          # for value residual learning
  
          orig_v = v
  
          if exists(value_residual):
              v = 0.5 * (v + value_residual)
  
          # kv cache
  
          if exists(kv_cache):
              ck, cv = kv_cache
              k = torch.cat((ck, k), dim = -2)
              v = torch.cat((cv, v), dim = -2)
  
          # store kv cache
  
          if return_kv_cache:
              kv_cache = torch.stack((k, v))
  
          # null key / values
  
          if self.num_null_kv > 0:
              null_k, null_v = repeat(self.null_kv, 'kv n d -> kv b n d', b = b).unbind(dim = 0)
              k = torch.cat((null_k, k), dim = -2)
              v = torch.cat((null_v, v), dim = -2)
  
          # split for multi-headed attention
  
          q = rearrange(q, 'b n (h d) -> b h n d', h = self.heads)
  
          # handle mask and null key / value
  
          if exists(mask):
              mask = F.pad(mask, (self.num_null_kv, 0), value = True)
  
          # attention
  
          out = self.attend(q, k, v, attn_bias = attn_bias, mask = mask)
  
          # merge heads
  
          out = rearrange(out, 'b h n d -> b n (h d)')
          out = self.to_out(out)
  
          if not return_kv_cache and not return_values:
              return out
  
          if return_kv_cache and not return_values:
              return out, kv_cache
  
          if return_values and not return_kv_cache:
              return out, orig_v
  
          return out, (kv_cache, orig_v)
  
  # transformer
  
  class Transformer(nn.Module):
      def __init__(
          self,
          *,
          dim,
          depth,
          heads,
          dim_context = None,
          cross_attend = False,
          attn_dropout = 0.,
          ff_dropout = 0.,
          grad_shrink_alpha = 0.1,
          cond_as_self_attn_prefix = False,
          rel_pos_bias = True,
          flash_attn = False,
          add_value_residual = True,
          num_residual_streams = 4,
          **kwargs
      ):
          super().__init__()
          rel_pos_bias = rel_pos_bias and not flash_attn
  
          assert not (cross_attend and cond_as_self_attn_prefix)
  
          self.dim_context = default(dim_context, dim)
  
          self.cond_as_self_attn_prefix = cond_as_self_attn_prefix
  
          self.grad_shrink = partial(grad_shrink, alpha = grad_shrink_alpha)
  
          self.layers = nn.ModuleList([])
  
          self.rel_pos_bias = RelativePositionBias(dim = dim // 2, heads = heads) if rel_pos_bias else None
  
          # hyper connections
  
          init_hyper_conn, self.expand_streams, self.reduce_streams = get_init_and_expand_reduce_stream_functions(num_residual_streams, disable = num_residual_streams == 1)
  
          # layers
  
          for _ in range(depth):
              self.layers.append(nn.ModuleList([
                  init_hyper_conn(dim = dim, branch = Attention(dim = dim, heads = heads, dropout = attn_dropout, flash = flash_attn, causal = True, **kwargs)),
                  init_hyper_conn(dim = dim, branch = Attention(dim = dim, heads = heads, dropout = attn_dropout, dim_context = dim_context, flash = flash_attn, num_null_kv = 1, norm_context = True, **kwargs)) if cross_attend else None,
                  init_hyper_conn(dim = dim, branch = FeedForward(dim = dim, dropout = ff_dropout))
              ]))
  
          self.norm = LayerNorm(dim)
  
          self.add_value_residual = add_value_residual
  
      def forward(
          self,
          x,
          self_attn_mask = None,
          context = None,
          context_mask = None,
          attn_bias = None,
          return_kv_cache = False,
          kv_cache = None
      ):
          assert not (self.cond_as_self_attn_prefix and not exists(context))
          assert not (exists(context) and context.shape[-1] != self.dim_context), f'you had specified a conditioning dimension of {self.dim_context}, yet what was received by the transformer has dimension of {context.shape[-1]}'
  
          n, device = x.shape[1], x.device
  
          # from cogview paper, adopted by GLM 130B LLM, decreases likelihood of attention net instability
  
          x = self.grad_shrink(x)
  
          # turn off kv cache if using conditioning as self attention (as in valle), for now
  
          if self.cond_as_self_attn_prefix:
              kv_cache = None
  
          # handle kv cache
  
          new_kv_cache = []
  
          if exists(kv_cache):
              cache_len = kv_cache.shape[-2]
              kv_cache = iter(kv_cache)
          else:
              cache_len = 0
              kv_cache = iter([])
  
          x = x[:, cache_len:]
  
          # relative positional bias
  
          if exists(attn_bias):
              rel_pos_bias = attn_bias
          else:
              rel_pos_bias = maybe(self.rel_pos_bias)(n, n)
  
          if exists(rel_pos_bias):
              rel_pos_bias = rel_pos_bias[..., cache_len:, :]
  
          # self attention kwargs
  
          self_attn_kwargs = dict()
          if self.cond_as_self_attn_prefix:
              self_attn_kwargs = dict(
                  prefix_context = context,
                  prefix_context_mask = context_mask
              )
  
          # value residuals
  
          self_attn_value_residual = None
          cross_attn_value_residual = None
  
          # expand residual streams
  
          x = self.expand_streams(x)
  
          # transformer layers
  
          for attn, cross_attn, ff in self.layers:
  
              residual = x
  
              x, (layer_kv_cache, values) = attn(x, attn_bias = rel_pos_bias, mask = self_attn_mask, kv_cache = next(kv_cache, None), return_kv_cache = True, return_values = True, value_residual = self_attn_value_residual, **self_attn_kwargs)
  
              if self.add_value_residual:
                  self_attn_value_residual = default(self_attn_value_residual, values)
  
              new_kv_cache.append(layer_kv_cache)
  
              if exists(cross_attn):
                  assert exists(context)
  
                  x, values = cross_attn(x, context = context, mask = context_mask, return_values = True, value_residual = cross_attn_value_residual)
  
                  if self.add_value_residual:
                      cross_attn_value_residual = default(cross_attn_value_residual, values)
  
              x = ff(x)
  
          # reduce residual streams
  
          x = self.reduce_streams(x)
  
          # final norm
  
          x = self.norm(x)
  
          if not return_kv_cache:
              return x
  
          return x, torch.stack(new_kv_cache)
  
  # the three hierarchical transformers
  
  class SemanticTransformer(nn.Module):
      @beartype
      def __init__(
          self,
          *,
          dim,
          depth,
          num_semantic_tokens,
          heads = 8,
          attn_dropout = 0.,
          ff_dropout = 0.,
          t5_name = DEFAULT_T5_NAME,
          cond_dim = None,
          has_condition = False,
          audio_text_condition = False,
          cond_as_self_attn_prefix = False,
          cond_drop_prob = 0.5,
          grad_shrink_alpha = 0.1,
          rel_pos_bias = True,
          flash_attn = False,
          **kwargs
      ):
          super().__init__()
          rel_pos_bias = rel_pos_bias and not flash_attn
  
          self.num_semantic_tokens = num_semantic_tokens
  
          if audio_text_condition:
              has_condition = True
              cond_dim = default(cond_dim, dim)
  
          self.has_condition = has_condition
          self.embed_text = partial(t5_encode_text, name = t5_name)
          self.cond_drop_prob = cond_drop_prob
  
          self.start_token = nn.Parameter(torch.randn(dim))
  
          self.semantic_embedding = nn.Embedding(num_semantic_tokens + 1, dim)
          self.eos_id = num_semantic_tokens
  
          text_dim = default(cond_dim, get_encoded_dim(t5_name))
          self.proj_text_embed = nn.Linear(text_dim, dim, bias = False) if text_dim != dim else nn.Identity()
  
          self.transformer = Transformer(
              dim = dim,
              depth = depth,
              heads = heads,
              attn_dropout = attn_dropout,
              ff_dropout = ff_dropout,
              cross_attend = has_condition and not cond_as_self_attn_prefix,
              cond_as_self_attn_prefix = cond_as_self_attn_prefix,
              grad_shrink_alpha = grad_shrink_alpha,
              rel_pos_bias = rel_pos_bias,
              flash_attn = flash_attn,
              **kwargs
          )
  
          self.to_logits = nn.Linear(dim, num_semantic_tokens + 1)
  
      @property
      def device(self):
          return next(self.parameters()).device
  
      def load(self, path):
          # Return pkg so that if this function gets called from within a Trainer function call,
          # the trainer can also access the package loaded from the checkpoint.
          device = self.device
          path = Path(path)
          assert path.exists()
          pkg = torch.load(str(path), map_location = device)
          # check version
          if 'version' in pkg and version.parse(pkg['version']) < version.parse(__version__):
              print(f'model was trained on older version {pkg["version"]} of audiolm-pytorch')
          self.load_state_dict(pkg['model'])
          return pkg
  
      def forward_with_cond_scale(
          self,
          *args,
          cond_scale = 3,
          kv_cache = None,
          return_kv_cache = False,
          **kwargs
      ):
          kv_cache = iter(default(kv_cache, []))
          new_kv_caches = []
  
          logits, new_kv_cache = self.forward(*args, cond_drop_prob = 0., kv_cache = next(kv_cache, None), return_kv_cache = True, **kwargs)
          new_kv_caches.append(new_kv_cache)
  
          if cond_scale == 1 or not self.has_condition:
              if not return_kv_cache:
                  return logits
  
              return logits, torch.stack(new_kv_caches)
  
          null_logits, null_new_kv_cache = self.forward(*args, cond_drop_prob = 1., kv_cache = next(kv_cache, None), return_kv_cache = True, **kwargs)
          new_kv_caches.append(null_new_kv_cache)
  
          scaled_logits = null_logits + (logits - null_logits) * cond_scale
  
          if not return_kv_cache:
              return scaled_logits
  
          return scaled_logits, torch.stack(new_kv_caches)
  
      @beartype
      def forward(
          self,
          *,
          ids = None,
          return_loss = False,
          text: list[str] | None = None,
          text_embeds = None,
          self_attn_mask = None,
          cond_drop_prob = None,
          unique_consecutive = None,
          kv_cache = None,
          return_kv_cache = False
      ):
          device = self.device
  
          b = ids.shape[0]
  
          has_text = exists(text) or exists(text_embeds)
          assert not (self.has_condition ^ has_text)
  
          text_mask = None
          if not exists(text_embeds) and exists(text):
              with torch.inference_mode():
                  text_embeds = self.embed_text(text, output_device = device)
                  text_mask = torch.any(text_embeds != 0, dim = -1)
  
          if exists(text_embeds):
              text_embeds = self.proj_text_embed(text_embeds)
  
          cond_drop_prob = default(cond_drop_prob, self.cond_drop_prob)
  
          if exists(text_mask) and cond_drop_prob > 0:
              keep_mask = prob_mask_like((b,), 1 - cond_drop_prob, device = device)
              text_mask = rearrange(keep_mask, 'b -> b 1') & text_mask
  
          if return_loss:
              labels, ids = ids.clone(), ids[:, :-1]
  
          tokens = get_embeds(self.semantic_embedding, ids)
  
          start_tokens = repeat(self.start_token, 'd -> b 1 d', b = ids.shape[0])
  
          tokens = torch.cat((start_tokens, tokens), dim = 1)
  
          if exists(self_attn_mask):
              self_attn_mask = F.pad(self_attn_mask, (1, 0), value = True)
  
          tokens, kv_cache = self.transformer(tokens, context = text_embeds, self_attn_mask = self_attn_mask, context_mask = text_mask, kv_cache = kv_cache, return_kv_cache = True)
          logits = self.to_logits(tokens)
  
          if not return_kv_cache:
              return logits
  
          return logits, kv_cache
  
  class CoarseTransformer(nn.Module):
      @beartype
      def __init__(
          self,
          *,
          codebook_size,
          num_coarse_quantizers,
          dim,
          depth,
          num_semantic_tokens,
          heads = 8,
          attn_dropout = 0.,
          ff_dropout = 0.,
          t5_name = DEFAULT_T5_NAME,
          has_condition = False,
          cond_dim = None,
          audio_text_condition = False,
          cond_as_self_attn_prefix = False,
          cond_drop_prob = 0.5,
          grad_shrink_alpha = 0.1,
          project_semantic_logits = True,
          rel_pos_bias = True,
          flash_attn = False,
          **kwargs
      ):
          super().__init__()
          rel_pos_bias = rel_pos_bias and not flash_attn
  
          self.num_semantic_tokens = num_semantic_tokens
  
          if audio_text_condition:
              has_condition = True
              cond_dim = default(cond_dim, dim)
  
          self.has_condition = has_condition
          self.embed_text = partial(t5_encode_text, name = t5_name)
          self.cond_drop_prob = cond_drop_prob
  
          self.semantic_start_token = nn.Parameter(torch.randn(dim))
          self.coarse_start_token = nn.Parameter(torch.randn(dim))
  
          self.semantic_eos_id = num_semantic_tokens
          self.semantic_embedding = nn.Embedding(num_semantic_tokens + 1, dim)
  
          self.coarse_eos_id = codebook_size
          codebook_size_with_eos = codebook_size + 1
  
          self.coarse_embedding = nn.Embedding(num_coarse_quantizers * codebook_size_with_eos, dim)
          self.coarse_quantize_embedding = nn.Embedding(num_coarse_quantizers, dim)
  
          text_dim = default(cond_dim, get_encoded_dim(t5_name))
          self.proj_text_embed = nn.Linear(text_dim, dim, bias = False) if text_dim != dim else nn.Identity()
  
          self.cross_attn_bias = nn.Parameter(torch.zeros(heads, 1, 1)) if rel_pos_bias else None
  
          self.transformer = Transformer(
              dim = dim,
              depth = depth,
              heads = heads,
              attn_dropout = attn_dropout,
              ff_dropout = ff_dropout,
              cross_attend = has_condition and not cond_as_self_attn_prefix,
              cond_as_self_attn_prefix = cond_as_self_attn_prefix,
              grad_shrink_alpha = grad_shrink_alpha,
              rel_pos_bias = rel_pos_bias,
              flash_attn = flash_attn,
              **kwargs
          )
  
          self.codebook_size = codebook_size
          self.num_coarse_quantizers = num_coarse_quantizers
  
          self.to_semantic_logits = nn.Linear(dim, num_semantic_tokens + 1) if project_semantic_logits else None
          self.coarse_logit_weights = nn.Parameter(torch.randn(num_coarse_quantizers, codebook_size_with_eos, dim))
  
      @property
      def device(self):
          return next(self.parameters()).device
  
      def load(self, path):
          # Return pkg so that if this function gets called from within a Trainer function call,
          # the trainer can also access the package loaded from the checkpoint.
          device = self.device
          path = Path(path)
          assert path.exists()
          pkg = torch.load(str(path), map_location = device)
          # check version
          if 'version' in pkg and version.parse(pkg['version']) < version.parse(__version__):
              print(f'model was trained on older version {pkg["version"]} of audiolm-pytorch')
          self.load_state_dict(pkg['model'])
          return pkg
  
      def forward_with_cond_scale(
          self,
          *args,
          cond_scale = 3,
          return_kv_cache = False,
          kv_cache = None,
          embed_cache = None,
          **kwargs
      ):
          iter_kv_cache = iter(default(kv_cache, []))
          iter_embed_cache = iter(default(embed_cache, []))
          new_kv_caches = []
          new_embed_caches = []
  
          (semantic_logits, coarse_logits), (new_kv_cache, new_embed_cache) = self.forward(*args, cond_drop_prob = 0., return_cache = True, kv_cache = next(iter_kv_cache, None), embed_cache = next(iter_embed_cache, None), **kwargs)
          new_kv_caches.append(new_kv_cache)
          new_embed_caches.append(new_embed_cache)
  
          if cond_scale == 1 or not self.has_condition:
              if not return_kv_cache:
                  return semantic_logits, coarse_logits
  
              return (semantic_logits, coarse_logits), (torch.stack(new_kv_caches), torch.stack(new_embed_caches))
  
          (null_semantic_logits, null_coarse_logits), (null_new_kv_cache, null_new_embed_cache) = self.forward(*args, cond_drop_prob = 1., return_cache = True, kv_cache = next(iter_kv_cache, None), embed_cache = next(iter_embed_cache, None), **kwargs)
          new_kv_caches.append(null_new_kv_cache)
          new_embed_caches.append(null_new_embed_cache)
  
          scaled_semantic_logits = None
          if exists(null_semantic_logits):
              scaled_semantic_logits = null_semantic_logits + (semantic_logits - null_semantic_logits) * cond_scale
  
          scaled_coarse_logits = null_coarse_logits + (coarse_logits - null_coarse_logits) * cond_scale
  
          if not return_kv_cache:
              return scaled_semantic_logits, scaled_coarse_logits
  
          return (scaled_semantic_logits, scaled_coarse_logits), (torch.stack(new_kv_caches), torch.stack(new_embed_caches))
  
      @beartype
      def forward(
          self,
          *,
          semantic_token_ids,
          coarse_token_ids,
          self_attn_mask = None,
          text: list[str] | None = None,
          text_embeds = None,
          cond_drop_prob = None,
          return_only_coarse_logits = False,
          return_cache = False,
          kv_cache = None,
          embed_cache = None
      ):
          b, device = semantic_token_ids.shape[0], semantic_token_ids.device
          arange = partial(torch.arange, device = device)
  
          has_text = exists(text) or exists(text_embeds)
          assert not (self.has_condition ^ has_text)
  
          if not exists(text_embeds) and exists(text):
              with torch.inference_mode():
                  text_embeds = self.embed_text(text, output_device = device)
  
          text_mask = None
          if exists(text_embeds):
              text_mask = torch.any(text_embeds != 0, dim = -1)
  
              text_embeds = self.proj_text_embed(text_embeds)
  
          cond_drop_prob = default(cond_drop_prob, self.cond_drop_prob)
  
          if exists(text_mask) and cond_drop_prob > 0:
              keep_mask = prob_mask_like((b,), 1 - cond_drop_prob, device = device)
              text_mask = rearrange(keep_mask, 'b -> b 1') & text_mask
  
          coarse_token_ids, semantic_token_ids = map(lambda t: rearrange(t, 'b ... -> b (...)'), (coarse_token_ids, semantic_token_ids))
  
          offsets = self.codebook_size * arange(self.num_coarse_quantizers)
          offsets = repeat(offsets, 'q -> 1 (n q)', n = ceil_div(coarse_token_ids.shape[-1], self.num_coarse_quantizers))
          offsets = offsets[:, :coarse_token_ids.shape[-1]]
          coarse_token_ids = coarse_token_ids + offsets
  
          semantic_tokens = get_embeds(self.semantic_embedding, semantic_token_ids)
          coarse_tokens = self.coarse_embedding(coarse_token_ids)
  
          coarse_quantize_tokens = repeat(self.coarse_quantize_embedding.weight, 'q d -> (n q) d', n = ceil_div(coarse_token_ids.shape[-1], self.num_coarse_quantizers))
          coarse_quantize_tokens = coarse_quantize_tokens[:coarse_token_ids.shape[-1], ...]
          coarse_tokens = coarse_tokens + coarse_quantize_tokens
  
          semantic_seq_len = semantic_tokens.shape[1]
  
          semantic_start_tokens = repeat(self.semantic_start_token, 'd -> b 1 d', b = b)
          coarse_start_tokens = repeat(self.coarse_start_token, 'd -> b 1 d', b = b)
  
          tokens = torch.cat((
              semantic_start_tokens,
              semantic_tokens,
              coarse_start_tokens,
              coarse_tokens
          ), dim = 1)
  
          # engineer the attention bias so that cross attention is not dominated by relative positions
  
          seq_len = tokens.shape[-2]
  
          attn_bias = None
  
          if exists(self.transformer.rel_pos_bias):
              attn_bias = self.transformer.rel_pos_bias(seq_len, seq_len)
  
              is_semantic = arange(seq_len) < (semantic_seq_len + 1) # semantic seq len + start token
              is_cross_attn = rearrange(is_semantic, 'i -> i 1') ^ rearrange(is_semantic, 'j -> 1 j')
  
              attn_bias = torch.where(
                  is_cross_attn,
                  self.cross_attn_bias,
                  attn_bias
              )
  
          # attend
  
          tokens, new_kv_cache = self.transformer(
              tokens,
              context = text_embeds,
              attn_bias = attn_bias,
              self_attn_mask = self_attn_mask,
              context_mask = text_mask,
              kv_cache = kv_cache,
              return_kv_cache = True
          )
  
          if exists(embed_cache):
              tokens = torch.cat((embed_cache, tokens), dim = -2)
  
          new_embed_cache = tokens
  
          # segment into semantic and coarse acoustic tokens
  
          pred_semantic_tokens, pred_coarse_tokens = tokens[:, :semantic_seq_len], tokens[:, (semantic_seq_len + 1):]
  
          # semantic logits
  
          semantic_logits = self.to_semantic_logits(pred_semantic_tokens) if not return_only_coarse_logits and exists(self.to_semantic_logits) else None
  
          # get coarse logits
  
          n = pred_coarse_tokens.shape[1]
          nq = round_down_nearest_multiple(n, self.num_coarse_quantizers)
  
          pred_coarse_tokens_groupable, pred_coarse_tokens_remainder = pred_coarse_tokens[:, :nq], pred_coarse_tokens[:, nq:]
  
          pred_coarse_tokens_groupable = rearrange(pred_coarse_tokens_groupable, 'b (n q) d -> b n q d', q = self.num_coarse_quantizers)
  
          coarse_logits_groupable = einsum('q c d, b n q d -> b n q c', self.coarse_logit_weights, pred_coarse_tokens_groupable)
  
          coarse_logits_groupable = rearrange(coarse_logits_groupable, 'b n q c -> b (n q) c')
  
          remainder_num_quantizers = pred_coarse_tokens_remainder.shape[1]
  
          if remainder_num_quantizers > 0:
              coarse_logits_remainder = einsum('q c d, b q d -> b q c', self.coarse_logit_weights[:remainder_num_quantizers], pred_coarse_tokens_remainder)
  
              coarse_logits = torch.cat((coarse_logits_groupable, coarse_logits_remainder), dim = 1)
          else:
              coarse_logits = coarse_logits_groupable
  
          logits = (semantic_logits, coarse_logits)
  
          if not return_cache:
              return logits
  
          return logits, (new_kv_cache, new_embed_cache)
  
  class FineTransformer(nn.Module):
      def __init__(
          self,
          *,
          num_coarse_quantizers,
          num_fine_quantizers,
          codebook_size,
          dim,
          depth,
          heads = 8,
          attn_dropout = 0.,
          ff_dropout = 0.,
          t5_name = DEFAULT_T5_NAME,
          has_condition = False,
          cond_dim = None,
          audio_text_condition = False,
          cond_as_self_attn_prefix = False,
          cond_drop_prob = 0.5,
          grad_shrink_alpha = 0.1,
          project_coarse_logits = True,
          pad_id = -1,
          rel_pos_bias = True,
          flash_attn = False,
          **kwargs
      ):
          super().__init__()
          rel_pos_bias = rel_pos_bias and not flash_attn
  
          if audio_text_condition:
              has_condition = True
              cond_dim = default(cond_dim, dim)
  
          self.has_condition = has_condition
          self.embed_text = partial(t5_encode_text, name = t5_name)
          self.cond_drop_prob = cond_drop_prob
  
  