def write_batch_samples_with_audio(
    batch: Batch,
    step: int,
    executor: concurrent.futures.Executor,
    result_dir: str,
    prefix: str,
    label_names: Optional[Dict[int, str]],
    config: ml_collections.ConfigDict,
):
  """Write a batch of samples to disk."""
  enable_condition = config.eval.get('results_with_condition', True)
  enable_target = config.eval.get('results_with_target', True)
  batch = apply_batch_mask(batch)
  label_ids = batch.get('label')
  if label_ids is not None and label_ids.shape[1] > 1:
    # Take the first one hot label
    label_ids = label_ids.argmax(axis=1, keepdims=True).astype(np.int32)
  idx_offset = step * config.eval_batch_size + jax.process_index(
  ) * config.eval_batch_size // jax.process_count()
  if 'frame_rate' in config.dataset_configs:
    fps = config.dataset_configs.frame_rate / config.dataset_configs.stride
  else:
    fps = 8
  audio_sample_rate = config.dataset_configs.audio_sample_rate
  batch_size = jax.tree_util.tree_leaves(batch)[0].shape[0]
  default_audio = batch.get('condition_audio', [None] * batch_size)
  dummy_video = np.zeros(
      (batch_size, config.dataset_configs.num_frames, 16, 16, 3), dtype=np.uint8
  )
  default_video = batch.get('condition_video', dummy_video)
  batch = batch.copy()  # shallow copy to fill in default values
  for key in ['condition', 'generated', 'target']:
    batch.setdefault(f'{key}_audio', default_audio)
    batch.setdefault(f'{key}_video', default_video)
  write_fn = functools.partial(executor.submit, write_video_with_audio, fps=fps,
                               audio_sample_rate=audio_sample_rate)
  for idx in range(batch_size):
    label_str = ''
    if label_names is not None and label_ids is not None:
      label_id = int(label_ids[idx, 0])
      label_str = f'{label_id:04d}{label_names[label_id]}_'
    filename = f'{prefix}_{label_str}{(idx + idx_offset):09d}'
    path = os.path.join(result_dir, filename)
    write_fn(f'{path}_gen.mp4', batch['generated_video'][idx],
             batch['generated_audio'][idx])
    if enable_condition:
      write_fn(f'{path}_cond.mp4', batch['condition_video'][idx],
               batch['condition_audio'][idx])
    if enable_target:
      write_fn(f'{path}_target.mp4', batch['target_video'][idx],
               batch['target_audio'][idx])
    if 'generated_video_sr' in batch:
      write_fn(
          f'{path}_gen_sr.mp4',
          batch['generated_video_sr'][idx],
          batch['generated_audio'][idx],
      )
Write a batch of samples to disk.
videogvt/train_lib/eval_utils.py