def get_vq_model(vq_config: ml_collections.ConfigDict) -> Dict[str, nn.Module]:
  """Build VQGAN generator and discriminator models."""
  config = vq_config
  if not ('vqgan' in config or 'model_class' in config or
          'model_type' in config.vqgan):
    raise ValueError('No complete vqgan found in the config.')
  model_info = f'{config.model_class}/{config.vqgan.model_type}'
  if model_info not in VQ_MODEL_INFO:
    raise ValueError(f'{model_info} not known.')
  model_dict = VQ_MODEL_INFO[model_info]
  dtype = get_dtype(config)
  precision = get_lax_precision(config)
  generator = model_dict['generator_cls'](
      config=config, dtype=dtype, precision=precision)
  discriminator = model_dict['discriminator_cls'](config=config, dtype=dtype)
  return {'generator': generator, 'discriminator': discriminator}
Build VQGAN generator and discriminator models.
videogvt/train_lib/train_utils.py