def eval_step_get_features(
    outputs: Dict[str, Any], *,
    metric_params: Dict[str, Any],
    model_suffix_list: Iterable[str] = ('',),
    config: ml_collections.ConfigDict) -> EvalFeatureDict:
  """Extract features to compute eval metrics.

  Args:
    outputs: Dict of original video, generated video and batch_mask.
    metric_params: Params for metric models.
    model_suffix_list: Tuple of model suffixes, '' for the default and '_ema'
      for the ema model.
    config: Configurations of the experiment.

  Returns:
    Extracted metric features.
  """
  features = {'batch_mask': outputs['batch_mask']}
  size_suffixes = ['']
  if config.eval.get('image_resize') is not None:
    suffix = f'_resize{config.eval.image_resize}'
    for key in [*outputs.keys()]:
      if 'video' in key:
        outputs[f'{key}{suffix}'] = metrics_lib.resize_bilinear(
            outputs[key], (config.eval.image_resize, config.eval.image_resize))
    size_suffixes.append(suffix)
    if config.eval.get('resized_only'):
      # Only evaluate resized samples, not original ones.
      size_suffixes.remove('')
  if config.eval.get('enable_inception_score', False):
    for s1, s2 in itertools.product(model_suffix_list, size_suffixes):
      suffix = s1 + s2
      features[f'inception_feature{suffix}'] = inception_score.run_model(
          metric_params['inception_score'],
          outputs[f'generated_video{suffix}'])
  if config.eval.get('enable_frechet_distance', False):
    for suffix in size_suffixes:
      if f'original_video{suffix}' not in outputs:
        continue
      features[f'frechet_feature_orig{suffix}'] = frechet_distance.run_model(
          metric_params['frechet_distance'], outputs[f'original_video{suffix}'],
          **config.eval.get('frechet_distance_args', {}))
    for s1, s2 in itertools.product(model_suffix_list, size_suffixes):
      suffix = s1 + s2
      if f'generated_video{suffix}' not in outputs:
        continue
      features[f'frechet_feature{suffix}'] = frechet_distance.run_model(
          metric_params['frechet_distance'],
          outputs[f'generated_video{suffix}'],
          **config.eval.get('frechet_distance_args', {}))
  if config.eval.get('enable_ssim_psnr', False):
    for s1, s2 in itertools.product(model_suffix_list, size_suffixes):
      suffix = s1 + s2
      cur_out_dict = image_quality_metrics.run_models(
          outputs[f'original_video{s2}'],
          outputs[f'generated_video{suffix}'],
          is_tf_function=False,
          metric_functions=None)
      features.update(
          {f'{k}{suffix}': v for k, v in cur_out_dict.items()})
      cur_out_dict = image_quality_metrics.run_models(
          outputs[f'generated_video{suffix}'][:, 1:],
          outputs[f'generated_video{suffix}'][:, :-1],
          is_tf_function=False,
          metric_functions=None)
      features.update(
          {f'interframe_{k}{suffix}': v for k, v in cur_out_dict.items()})
  if config.eval.get('enable_utilization', False):
    for suffix in model_suffix_list:
      if f'generated_tokens{suffix}' not in outputs:
        continue
      tokens = outputs[f'generated_tokens{suffix}']
      tokens = jax.tree_util.tree_map(
          lambda x: x.reshape(x.shape[0], -1), tokens)
      tokens = jnp.concatenate(jax.tree_leaves(tokens), axis=-1)
      features[f'tokens{suffix}'] = tokens

  features = jax.lax.all_gather(features, axis_name='device', tiled=True)
  return features  # total_batch_size, ...
Extract features to compute eval metrics.

  Args:
    outputs: Dict of original video, generated video and batch_mask.
    metric_params: Params for metric models.
    model_suffix_list: Tuple of model suffixes, '' for the default and '_ema'
      for the ema model.
    config: Configurations of the experiment.

  Returns:
    Extracted metric features.
  
videogvt/train_lib/eval_utils.py