text: 
  from wavegrad.preprocess import transform
  import numpy as np
  import os
  import torch
  import torchaudio
  
  from argparse import ArgumentParser
  
  from wavegrad.params import AttrDict, params as base_params
  from wavegrad.model import WaveGrad
  
  from time import time
  from armory.data.utils import maybe_download_weights_from_s3
  
  models = {}
  
  
  class Denoiser():
    def __init__(self, model_dir, params, device=torch.device('cuda')):
      # Lazy load model.
      if not model_dir in models:
        if os.path.exists(f'{model_dir}/weights.pt'):
          checkpoint = torch.load(f'{model_dir}/weights.pt')
        else:
          weights_path = maybe_download_weights_from_s3(model_dir)
          checkpoint = torch.load(weights_path)
          # checkpoint = torch.load(model_dir)
        model = WaveGrad(AttrDict(base_params)).to(device)
        model.load_state_dict(checkpoint['model'])
        model.eval()
        models[model_dir] = model
  
      model = models[model_dir]
      model.params.override(params)
  
      beta = np.array(model.params.noise_schedule)
      alpha = 1 - beta
      alpha_cum = np.cumprod(alpha)
  
      self.alpha = alpha
      self.alpha_cum = alpha_cum
      self.beta = beta
      self.model = model
      self.device = device
  
  
    def __call__(self, spectrogram, audio, severity=None, pin=None, output_name="output"):
      if severity is None:
        severity = len(self.model.params.noise_schedule)
  
      with torch.no_grad():
        # Expand rank 2 tensors by adding a batch dimension.
        if len(spectrogram.shape) == 2:
          spectrogram = spectrogram.unsqueeze(0)
        
        spectrogram = spectrogram.to(self.device)
  
        length = self.model.params.hop_samples * spectrogram.shape[-1]
  
        # if audio is None:
        #   audio = torch.randn(spectrogram.shape[0], length, device=self.device)
        # else:
        # TODO FROME HERE: padding or truncation
        if len(audio.shape) == 1:
          audio = audio.unsqueeze(0)
  
        if audio.shape[-1] >= length:
          audio = audio[..., :length]
        else:
          padding = (
            torch.zeros([audio.shape[0], length - audio.shape[1]])
            .to(audio.device)
          )
          audio = torch.cat([audio, padding], -1)
        
        audio = audio.to(self.device)
  
        noise_scale = (
          torch.from_numpy(self.alpha_cum**0.5)
          .float()
          .unsqueeze(1)
          .to(self.device)
        )
  
        ti = time()
        for n in range(severity - 1, -1, -1):
          # ===================
          if pin is not None:
            n = pin
          # ===================
  
          # print(f"{n:>3d}/{len(self.alpha):>3d}", end="\r")
  
          c1 = 1 / self.alpha[n]**0.5
          c2 = (1 - self.alpha[n]) / (1 - self.alpha_cum[n])**0.5
          prediction = self.model(audio, spectrogram, noise_scale[n]).squeeze(1)
          audio = c1 * (audio - c2 * prediction)
          if n > 0:
            noise = torch.randn_like(audio).to(self.device)
            sigma = ((1.0 - self.alpha_cum[n-1]) / (1.0 - self.alpha_cum[n]) * self.beta[n])**0.5
            audio += sigma * noise
          audio = torch.clamp(audio, -1.0, 1.0)   # TODO: J: I disagree with this step
  
      # print(f"Finished {spectrogram.shape} in {time() - ti:.2f} secs.", end="\r")
      return audio
  
  
  
  
  
  def main(args):
    # audio = torchaudio.load_wav(args.audio)
    # spectrogram = torch.from_numpy(np.load(args.spectrogram_path))
    
    spectrogram, audio = transform(args.audio_path)
  
    params = {}
    if args.noise_schedule:
      params['noise_schedule'] = torch.from_numpy(np.load(args.noise_schedule))
  
    denoiser = Denoiser(args.model_dir, params)
  
    for s in args.severity:
      if s is None:
        s = len(denoiser.alpha)
      audio = denoiser(spectrogram, audio, severity=s)
      torchaudio.save(
        f"{args.output}-severity-{s:03d}.wav",
        audio.cpu(),
        sample_rate=denoiser.model.params.sample_rate)
  
  
  if __name__ == '__main__':
    parser = ArgumentParser(
        description='runs inference on a spectrogram file generated by wavegrad.preprocess')
    parser.add_argument(
      'model_dir',
      help='directory containing a trained model (or full path to weights.pt file)')
    parser.add_argument(
      'audio_path',
      help='path to a audio file generated by wavegrad.preprocess')
    parser.add_argument(
      '--noise-schedule', '-n', default=None,
      help='path to a custom noise schedule file generated by wavegrad.noise_schedule')
    parser.add_argument(
      '--output', '-o', default='output',
      help='output file name')
    parser.add_argument("--severity", nargs="*", type=int, default=[None])
    parser.add_argument("--pin", type=int, default=None)
    main(parser.parse_args())
  