Dynamic caption: How to use a dedicated caption field on a media entity?

Right now you can't use media tokens when using photoswipe field formatter on node for media entity reference field. That's because this code $entity_type = $imageDTO->getEntity()->getEntityTypeId();(source) returns a host entity (node) when formatter being called from node rendering.
I have a similar situation but we are using paragraph for storing image with additional data. We're also using custom caption field but inside paragraph. I modified https://git.drupalcode.org/project/photoswipe/-/blob/5.x/modules/photosw... by this:

case 'custom':
        $entity_type = $imageDTO->getEntity()->getEntityTypeId();
        $entity = $imageDTO->getEntity();
        $token = $thirdPartySettings['photoswipe_dynamic_caption']['photoswipe_caption_custom'];

        // Try to extract parent entity type from token.
        $token_parts = explode(':', str_replace(['[', ']'], '', $token));
        $token_entity_type = $token_parts[0] ?? NULL;
        // Let's use parent entity for token generation.
        if (count($token_parts) > 1) {
          $entity_type = $token_entity_type;
          $entity = $imageDTO->getItem()->getParent()->getEntity();
        }

        $caption = \Drupal::service('token')->replace($token,
          [
            $entity_type => $entity,
            'file' => $imageDTO->getItem(),
          ],
          [
            'clear' => TRUE,
            'langcode' => \Drupal::service('language_manager')->getCurrentLanguage()->getId(),
          ]
        );
        break;

I agree it would be helpful to add related / parent entity tokens, but I dislike this approach:

        // Try to extract parent entity type from token.
        $token_parts = explode(':', str_replace(['[', ']'], '', $token));
        $token_entity_type = $token_parts[0] ?? NULL;

Any other ideas how to add the references token more cleanly? Has this perhaps already been solved somewhere else?