code: def predict_mel(tokens, durations, ckpt_fn):
      global dic
      with open(ckpt_fn, "rb") as f:
          if dic is None:
              dic = pickle.load(f)
          last_step, params, aux, rng, optim_state = (
              dic["step"],
              dic["params"],
              dic["aux"],
              dic["rng"],
              dic["optim_state"],
          )
  
      @hk.transform_with_state
      def forward(tokens, durations, n_frames):
          net = AcousticModel(is_training=False)
          return net.inference(tokens, durations, n_frames)
  
      durations = durations * FLAGS.sample_rate / (FLAGS.n_fft // 4)
      n_frames = int(jnp.sum(durations).item())
      predict_fn = jax.jit(forward.apply, static_argnums=[5])
      tokens = np.array(tokens, dtype=np.int32)[None, :]
      return predict_fn(params, aux, rng, tokens, durations, n_frames)[
          0
      ]
  
file_path: underthesea/underthesea/pipeline/say/viettts_/nat/text2mel.py
imports: import pickle
  from argparse import ArgumentParser
  from pathlib import Path
  import haiku
  import jax
  import jax.numpy
  import matplotlib.pyplot
  import numpy
  from .config import FLAGS, DurationInput
  from .data_loader import load_phonemes_set
  from .model import AcousticModel, DurationModel
label_desc: Use image node_id 4 for calling a global function with example usage: predict_mel(tokens, durations, ckpt_fn) and returns: predict_fn, net