def get_task_registry(config: ml_collections.ConfigDict,
                      is_train: bool = True) -> Dict[str, Task]:
  """Get the mapping from task name to task."""
  task_ids = config.get('tasks', ('full_generation',))
  tasks = []

  for task_id in task_ids:
    if task_id.startswith('full_generation'):
      tasks.append(get_full_generation_task(config, is_train, task_id))
    elif task_id.startswith('latent_frame_prediction'):
      tasks.append(get_latent_frame_prediction_task(config, is_train, task_id))
    elif task_id.startswith('frame_prediction'):
      tasks.append(get_frame_prediction_task(config, is_train, task_id))
    elif task_id.startswith('frame_interpolation'):
      tasks.append(get_frame_interpolation_task(config, is_train, task_id))
    elif task_id.startswith('outpainting'):
      tasks.append(get_outpainting_task(config, is_train, task_id))
    elif task_id.startswith('inpainting'):
      tasks.append(get_inpainting_task(config, is_train, task_id))
    else:
      raise ValueError(f'Unknown task_id: {task_id}')
  tasks = {task.task_id: task for task in tasks}
  return tasks
Get the mapping from task name to task.
videogvt/train_lib/task_registry.py