def save_img_uint8(img, pth):
  """Save an image (probably RGB) in [0, 1] to disk as a uint8 PNG."""
  with open_file(pth, 'wb') as f:
    Image.fromarray(
        (np.clip(np.nan_to_num(img), 0., 1.) * 255.).astype(jnp.uint8)).save(
            f, 'PNG')
Save an image (probably RGB) in [0, 1] to disk as a uint8 PNG.
internal/utils.py