diff --git a/src/EntityEmbedDisplay/EntityEmbedDisplayBase.php b/src/EntityEmbedDisplay/EntityEmbedDisplayBase.php
index b9c5c19..3c55f0b 100644
--- a/src/EntityEmbedDisplay/EntityEmbedDisplayBase.php
+++ b/src/EntityEmbedDisplay/EntityEmbedDisplayBase.php
@@ -4,13 +4,14 @@ namespace Drupal\entity_embed\EntityEmbedDisplay;
 
 use Drupal\Component\Utility\NestedArray;
 use Drupal\Core\Access\AccessResult;
-use Drupal\Core\Entity\EntityManagerInterface;
+use Drupal\Core\Entity\EntityTypeManagerInterface;
 use Drupal\Core\Form\FormStateInterface;
 use Drupal\Core\Language\LanguageInterface;
+use Drupal\Core\Language\LanguageManagerInterface;
 use Drupal\Core\Plugin\ContainerFactoryPluginInterface;
 use Drupal\Core\Plugin\PluginBase;
 use Drupal\Core\Session\AccountInterface;
-use Drupal\entity_embed\EntityHelperTrait;
+use Drupal\entity_embed\EntityEmbedHelperTrait;
 use Symfony\Component\DependencyInjection\ContainerInterface;
 
 /**
@@ -24,7 +25,21 @@ use Symfony\Component\DependencyInjection\ContainerInterface;
   * @ingroup entity_embed_api
  */
 abstract class EntityEmbedDisplayBase extends PluginBase implements ContainerFactoryPluginInterface, EntityEmbedDisplayInterface {
-  use EntityHelperTrait;
+  use EntityEmbedHelperTrait;
+
+  /**
+   * The entity type manager service.
+   *
+   * @var \Drupal\Core\Entity\EntityTypeManagerInterface
+   */
+  protected $entityTypeManager;
+
+  /**
+   * The language manager.
+   *
+   * @var \Drupal\Core\Language\LanguageManagerInterface $language_manager
+   */
+  protected $languageManager;
 
   /**
    * The context for the plugin.
@@ -43,13 +58,16 @@ abstract class EntityEmbedDisplayBase extends PluginBase implements ContainerFac
   /**
    * {@inheritdoc}
    *
-   * @param \Drupal\Core\Entity\EntityManagerInterface $entity_manager
-   *   The entity manager service.
+   * @param \Drupal\Core\Entity\EntityTypeManagerInterface $entity_type_manager
+   *   The entity type manager service.
+   * @param \Drupal\Core\Language\LanguageManagerInterface $language_manager
+   *   The language manager.
    */
-  public function __construct(array $configuration, $plugin_id, $plugin_definition, EntityManagerInterface $entity_manager) {
+  public function __construct(array $configuration, $plugin_id, $plugin_definition, EntityTypeManagerInterface $entity_type_manager, LanguageManagerInterface $language_manager) {
     parent::__construct($configuration, $plugin_id, $plugin_definition);
     $this->setConfiguration($configuration);
-    $this->setEntityManager($entity_manager);
+    $this->entityTypeManager = $entity_type_manager;
+    $this->languageManager = $language_manager;
   }
 
   /**
@@ -60,7 +78,8 @@ abstract class EntityEmbedDisplayBase extends PluginBase implements ContainerFac
       $configuration,
       $plugin_id,
       $plugin_definition,
-      $container->get('entity.manager')
+      $container->get('entity_type.manager'),
+      $container->get('language_manager')
     );
   }
 
@@ -92,7 +111,7 @@ abstract class EntityEmbedDisplayBase extends PluginBase implements ContainerFac
     // First, determine whether or not the entity type id is valid. Return FALSE
     // if the specified id is not valid.
     $entity_type = $this->getEntityTypeFromContext();
-    if (!$this->entityManager()->getDefinition($entity_type)) {
+    if (!$this->entityTypeManager->getDefinition($entity_type)) {
       return FALSE;
     }
 
@@ -308,8 +327,9 @@ abstract class EntityEmbedDisplayBase extends PluginBase implements ContainerFac
   public function getLangcode() {
     $langcode = $this->getAttributeValue('data-langcode');
     if (empty($langcode)) {
-      $langcode = \Drupal::languageManager()->getCurrentLanguage(LanguageInterface::TYPE_CONTENT)->getId();
+      $langcode = $this->languageManager->getCurrentLanguage(LanguageInterface::TYPE_CONTENT)->getId();
     }
     return $langcode;
   }
+
 }
diff --git a/src/EntityEmbedDisplay/FieldFormatterEntityEmbedDisplayBase.php b/src/EntityEmbedDisplay/FieldFormatterEntityEmbedDisplayBase.php
index 53aa02a..595fbda 100644
--- a/src/EntityEmbedDisplay/FieldFormatterEntityEmbedDisplayBase.php
+++ b/src/EntityEmbedDisplay/FieldFormatterEntityEmbedDisplayBase.php
@@ -3,10 +3,11 @@
 namespace Drupal\entity_embed\EntityEmbedDisplay;
 
 use Drupal\Core\Access\AccessResult;
-use Drupal\Core\Entity\EntityManagerInterface;
+use Drupal\Core\Entity\EntityTypeManagerInterface;
 use Drupal\Core\Field\BaseFieldDefinition;
 use Drupal\Core\Field\FormatterPluginManager;
 use Drupal\Core\Form\FormStateInterface;
+use Drupal\Core\Language\LanguageManagerInterface;
 use Drupal\Core\Plugin\PluginDependencyTrait;
 use Drupal\Core\Session\AccountInterface;
 use Drupal\Core\TypedData\TypedDataManager;
@@ -47,19 +48,20 @@ abstract class FieldFormatterEntityEmbedDisplayBase extends EntityEmbedDisplayBa
   /**
    * Constructs a FieldFormatterEntityEmbedDisplayBase object.
    *
-   * @param \Drupal\Core\Entity\EntityManagerInterface $entity_manager
-   *   The entity manager service.
+   * @param \Drupal\Core\Entity\EntityTypeManagerInterface $entity_type_manager
+   *   The entity type manager service.
    * @param \Drupal\Core\Field\FormatterPluginManager $formatter_plugin_manager
    *   The field formatter plugin manager.
    * @param \Drupal\Core\TypedData\TypedDataManager $typed_data_manager
    *   The typed data manager.
+   * @param \Drupal\Core\Language\LanguageManagerInterface $language_manager
+   *   The language manager.
    */
-  public function __construct(array $configuration, $plugin_id, $plugin_definition, EntityManagerInterface $entity_manager, FormatterPluginManager $formatter_plugin_manager, TypedDataManager $typed_data_manager) {
+  public function __construct(array $configuration, $plugin_id, $plugin_definition, EntityTypeManagerInterface $entity_type_manager, FormatterPluginManager $formatter_plugin_manager, TypedDataManager $typed_data_manager, LanguageManagerInterface $language_manager) {
     $this->formatterPluginManager = $formatter_plugin_manager;
     $this->setConfiguration($configuration);
-    $this->setEntityManager($entity_manager);
     $this->typedDataManager = $typed_data_manager;
-    parent::__construct($configuration, $plugin_id, $plugin_definition, $entity_manager);
+    parent::__construct($configuration, $plugin_id, $plugin_definition, $entity_type_manager, $language_manager);
   }
 
   /**
@@ -70,9 +72,10 @@ abstract class FieldFormatterEntityEmbedDisplayBase extends EntityEmbedDisplayBa
       $configuration,
       $plugin_id,
       $plugin_definition,
-      $container->get('entity.manager'),
+      $container->get('entity_type.manager'),
       $container->get('plugin.manager.field.formatter'),
-      $container->get('typed_data_manager')
+      $container->get('typed_data_manager'),
+      $container->get('language_manager')
     );
   }
 
diff --git a/src/EntityEmbedHelperTrait.php b/src/EntityEmbedHelperTrait.php
new file mode 100644
index 0000000..180dd36
--- /dev/null
+++ b/src/EntityEmbedHelperTrait.php
@@ -0,0 +1,131 @@
+<?php
+
+namespace Drupal\entity_embed;
+
+use Drupal\Component\Utility\Html;
+use Drupal\Core\Entity\EntityInterface;
+
+/**
+ * Wrapper methods for entity rendering.
+ *
+ * Trait depends on EntityEmbedDisplayManager and ModuleHandlerInterface being
+ * available as class members ($displayPluginManager and $moduleHandler). Any
+ * classes using this trait should make sure that both dependencies are
+ * provided.
+ *
+ * @internal
+ */
+trait EntityEmbedHelperTrait {
+
+  /**
+   * Builds the render array for an embedded entity.
+   *
+   * @param \Drupal\Core\Entity\EntityInterface $entity
+   *   The entity to be rendered.
+   * @param array $context
+   *   (optional) Array of context values, corresponding to the attributes on
+   *   the embed HTML tag.
+   *
+   * @return array
+   *   A render array.
+   *
+   * @todo improve documentation
+   */
+  protected function buildEntityEmbed(EntityInterface $entity, array $context = array()) {
+    // Support the deprecated view-mode data attribute.
+    if (isset($context['data-view-mode']) && !isset($context['data-entity-embed-display']) && !isset($context['data-entity-embed-settings'])) {
+      $context['data-entity-embed-display'] = 'entity_reference:entity_reference_entity_view';
+      $context['data-entity-embed-settings'] = ['view_mode' => &$context['data-view-mode']];
+    }
+
+    // Merge in default attributes.
+    $context += array(
+      'data-entity-type' => $entity->getEntityTypeId(),
+      'data-entity-uuid' => $entity->uuid(),
+      'data-entity-embed-display' => 'entity_reference:entity_reference_entity_view',
+      'data-entity-embed-settings' => array(),
+    );
+
+    // The default Entity Embed Display plugin has been deprecated by the
+    // rendered entity field formatter.
+    if ($context['data-entity-embed-display'] === 'default') {
+      $context['data-entity-embed-display'] = 'entity_reference:entity_reference_entity_view';
+    }
+
+    // The caption text is double-encoded, so decode it here.
+    if (isset($context['data-caption'])) {
+      $context['data-caption'] = Html::decodeEntities($context['data-caption']);
+    }
+
+    // Allow modules to alter the entity prior to embed rendering.
+    $this->moduleHandler->alter(array("{$context['data-entity-type']}_embed_context", 'entity_embed_context'), $context, $entity);
+
+    // Build and render the Entity Embed Display plugin, allowing modules to
+    // alter the result before rendering.
+    $build = array(
+      '#theme_wrappers' => ['entity_embed_container'],
+      '#attributes' => ['class' => ['embedded-entity']],
+      '#entity' => $entity,
+      '#context' => $context,
+    );
+    $build['entity'] = $this->buildEntityEmbedDisplayPlugin(
+      $entity,
+      $context['data-entity-embed-display'],
+      $context['data-entity-embed-settings'],
+      $context
+    );
+
+    // Maintain data-align if it is there.
+    if (isset($context['data-align'])) {
+      $build['#attributes']['data-align'] = $context['data-align'];
+    }
+    elseif ((isset($context['class']))) {
+      $build['#attributes']['class'][] = $context['class'];
+    }
+
+    // Maintain data-caption if it is there.
+    if (isset($context['data-caption'])) {
+      $build['#attributes']['data-caption'] = $context['data-caption'];
+    }
+
+    // Make sure that access to the entity is respected.
+    $build['#access'] = $entity->access('view', NULL, TRUE);
+
+    // @todo Should this hook get invoked if $build is an empty array?
+    $this->moduleHandler->alter(array("{$context['data-entity-type']}_embed", 'entity_embed'), $build, $entity, $context);
+    return $build;
+  }
+
+  /**
+   * Builds the render array for an entity using an Entity Embed Display plugin.
+   *
+   * @param \Drupal\Core\Entity\EntityInterface $entity
+   *   The entity to be rendered.
+   * @param string $plugin_id
+   *   The Entity Embed Display plugin ID.
+   * @param array $plugin_configuration
+   *   (optional) Array of plugin configuration values.
+   * @param array $context
+   *   (optional) Array of additional context values, usually the embed HTML
+   *   tag's attributes.
+   *
+   * @return array
+   *   A render array for the Entity Embed Display plugin.
+   */
+  protected function buildEntityEmbedDisplayPlugin(EntityInterface $entity, $plugin_id, array $plugin_configuration = array(), array $context = array()) {
+    // Build the Entity Embed Display plugin.
+    /** @var \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayBase $display */
+    $display = $this->displayPluginManager->createInstance($plugin_id, $plugin_configuration);
+    $display->setContextValue('entity', $entity);
+    $display->setAttributes($context);
+
+    // Check if the Entity Embed Display plugin is accessible. This also checks
+    // entity access, which is why we never call $entity->access() here.
+    if (!$display->access()) {
+      return array();
+    }
+
+    return $display->build();
+  }
+
+}
diff --git a/src/EntityHelperTrait.php b/src/EntityHelperTrait.php
deleted file mode 100644
index 1516507..0000000
--- a/src/EntityHelperTrait.php
+++ /dev/null
@@ -1,339 +0,0 @@
-<?php
-
-namespace Drupal\entity_embed;
-
-use Drupal\Component\Utility\Html;
-use Drupal\Core\Entity\EntityInterface;
-use Drupal\Core\Entity\EntityManagerInterface;
-use Drupal\Core\Entity\EntityStorageException;
-use Drupal\Core\Extension\ModuleHandlerInterface;
-use Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager;
-
-/**
- * Wrapper methods for entity loading and rendering.
- *
- * This utility trait should only be used in application-level code, such as
- * classes that would implement ContainerInjectionInterface. Services registered
- * in the Container should not use this trait but inject the appropriate service
- * directly for easier testing.
- *
- * @todo this duplicates/wraps much of the Entity API. Is this really worth
- * keeping? The downside is painfully illustrated: the documentation must be
- * kept up to date with the actual API documentation… and that's not happening.
- * This causes subtle bugs and makes maintenance harder.
- */
-trait EntityHelperTrait {
-
-  /**
-   * The entity manager service.
-   *
-   * @var \Drupal\Core\Entity\EntityManagerInterface
-   */
-  protected $entityManager;
-
-  /**
-   * The module handler service.
-   *
-   * @var \Drupal\Core\Extension\ModuleHandlerInterface.
-   */
-  protected $moduleHandler;
-
-  /**
-   * The Entity Embed Display plugin manager.
-   *
-   * @var \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager.
-   */
-  protected $displayPluginManager;
-
-  /**
-   * Loads an entity from the database.
-   *
-   * @param string $entity_type
-   *   The entity type to load, e.g. node or user.
-   * @param mixed $id
-   *   The id or UUID of the entity to load.
-   *
-   * @return \Drupal\Core\Entity\EntityInterface
-   *   The entity object, or NULL if there is no entity with the given id or
-   *   UUID.
-   */
-  protected function loadEntity($entity_type, $id) {
-    $entities = $this->loadMultipleEntities($entity_type, array($id));
-    return !empty($entities) ? reset($entities) : NULL;
-  }
-
-  /**
-   * Loads multiple entities from the database.
-   *
-   * @param string $entity_type
-   *   The entity type to load, e.g. node or user.
-   * @param array $ids
-   *   An array of entity IDs or UUIDs.
-   *
-   * @return array
-   *   An array of entity objects indexed by their ids.
-   *
-   * @throws \Drupal\Core\Entity\EntityStorageException
-   *   Throws an exception if the entity type does not supports UUIDs.
-   */
-  protected function loadMultipleEntities($entity_type, array $ids) {
-    $entities = array();
-    $storage = $this->entityManager()->getStorage($entity_type);
-
-    $uuids = array_filter($ids, 'Drupal\Component\Uuid\Uuid::isValid');
-    if (!empty($uuids)) {
-      $definition = $this->entityManager()->getDefinition($entity_type);
-      if (!$uuid_key = $definition->getKey('uuid')) {
-        throw new EntityStorageException("Entity type $entity_type does not support UUIDs.");
-      }
-      $entities += $storage->loadByProperties(array($uuid_key => $uuids));
-    }
-
-    if ($remaining_ids = array_diff($ids, $uuids)) {
-      $entities += $storage->loadMultiple($remaining_ids);
-    }
-
-    return $entities;
-  }
-
-  /**
-   * Determines if an entity can be rendered.
-   *
-   * @param \Drupal\Core\Entity\EntityInterface $entity
-   *   The entity object.
-   *
-   * @return bool
-   *   TRUE if the entity's type has a view builder controller, otherwise FALSE.
-   */
-  protected function canRenderEntity(EntityInterface $entity) {
-    $entity_type = $entity->getEntityTypeId();
-    return $this->canRenderEntityType($entity_type);
-  }
-
-  /**
-   * Determines if an entity type can be rendered.
-   *
-   * @param string $entity_type
-   *   The entity type id.
-   *
-   * @return bool
-   *   TRUE if the entitys type has a view builder controller, otherwise FALSE.
-   */
-  protected function canRenderEntityType($entity_type) {
-    return $this->entityManager()->hasHandler($entity_type, 'view_builder');
-  }
-
-  /**
-   * Builds the render array for the provided entity.
-   *
-   * @param \Drupal\Core\Entity\EntityInterface $entity
-   *   The entity to be rendered.
-   * @param string $view_mode
-   *   The view mode that should be used to display the entity.
-   * @param string $langcode
-   *   (optional) For which language the entity should be rendered, defaults to
-   *   the current content language.
-   *
-   * @return array
-   *   A render array for the entity.
-   *
-   * @see \Drupal\Core\Entity\EntityViewBuilderInterface::view()
-   *
-   * @todo Note that the signature here does NOT match that of \Drupal\Core\Entity\EntityViewBuilderInterface::view(), which can lead to subtle bugs.
-   */
-  protected function buildEntity(EntityInterface $entity, $view_mode, $langcode = NULL) {
-    return $this->entityManager()
-      ->getViewBuilder($entity->getEntityTypeId())
-      ->view($entity, $view_mode, $langcode);
-  }
-
-  /**
-   * Builds the render array for an embedded entity.
-   *
-   * @param \Drupal\Core\Entity\EntityInterface $entity
-   *   The entity to be rendered.
-   * @param array $context
-   *   (optional) Array of context values, corresponding to the attributes on
-   *   the embed HTML tag.
-   *
-   * @return array
-   *   A render array.
-   *
-   * @todo improve documentation
-   */
-  protected function buildEntityEmbed(EntityInterface $entity, array $context = array()) {
-    // Support the deprecated view-mode data attribute.
-    if (isset($context['data-view-mode']) && !isset($context['data-entity-embed-display']) && !isset($context['data-entity-embed-settings'])) {
-      $context['data-entity-embed-display'] = 'entity_reference:entity_reference_entity_view';
-      $context['data-entity-embed-settings'] = ['view_mode' => &$context['data-view-mode']];
-    }
-
-    // Merge in default attributes.
-    $context += array(
-      'data-entity-type' => $entity->getEntityTypeId(),
-      'data-entity-uuid' => $entity->uuid(),
-      'data-entity-embed-display' => 'entity_reference:entity_reference_entity_view',
-      'data-entity-embed-settings' => array(),
-    );
-
-    // The default Entity Embed Display plugin has been deprecated by the
-    // rendered entity field formatter.
-    if ($context['data-entity-embed-display'] === 'default') {
-      $context['data-entity-embed-display'] = 'entity_reference:entity_reference_entity_view';
-    }
-
-    // The caption text is double-encoded, so decode it here.
-    if (isset($context['data-caption'])) {
-      $context['data-caption'] = Html::decodeEntities($context['data-caption']);
-    }
-
-    // Allow modules to alter the entity prior to embed rendering.
-    $this->moduleHandler()->alter(array("{$context['data-entity-type']}_embed_context", 'entity_embed_context'), $context, $entity);
-
-    // Build and render the Entity Embed Display plugin, allowing modules to
-    // alter the result before rendering.
-    $build = array(
-      '#theme_wrappers' => ['entity_embed_container'],
-      '#attributes' => ['class' => ['embedded-entity']],
-      '#entity' => $entity,
-      '#context' => $context,
-    );
-    $build['entity'] = $this->buildEntityEmbedDisplayPlugin(
-      $entity,
-      $context['data-entity-embed-display'],
-      $context['data-entity-embed-settings'],
-      $context
-    );
-
-    // Maintain data-align if it is there.
-    if (isset($context['data-align'])) {
-      $build['#attributes']['data-align'] = $context['data-align'];
-    }
-    elseif ((isset($context['class']))) {
-      $build['#attributes']['class'][] = $context['class'];
-    }
-
-    // Maintain data-caption if it is there.
-    if (isset($context['data-caption'])) {
-      $build['#attributes']['data-caption'] = $context['data-caption'];
-    }
-
-    // Make sure that access to the entity is respected.
-    $build['#access'] = $entity->access('view', NULL, TRUE);
-
-    // @todo Should this hook get invoked if $build is an empty array?
-    $this->moduleHandler()->alter(array("{$context['data-entity-type']}_embed", 'entity_embed'), $build, $entity, $context);
-    return $build;
-  }
-
-  /**
-   * Builds the render array for an entity using an Entity Embed Display plugin.
-   *
-   * @param \Drupal\Core\Entity\EntityInterface $entity
-   *   The entity to be rendered.
-   * @param string $plugin_id
-   *   The Entity Embed Display plugin ID.
-   * @param array $plugin_configuration
-   *   (optional) Array of plugin configuration values.
-   * @param array $context
-   *   (optional) Array of additional context values, usually the embed HTML
-   *   tag's attributes.
-   *
-   * @return array
-   *   A render array for the Entity Embed Display plugin.
-   */
-  protected function buildEntityEmbedDisplayPlugin(EntityInterface $entity, $plugin_id, array $plugin_configuration = array(), array $context = array()) {
-    // Build the Entity Embed Display plugin.
-    /** @var \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayBase $display */
-    $display = $this->displayPluginManager()->createInstance($plugin_id, $plugin_configuration);
-    $display->setContextValue('entity', $entity);
-    $display->setAttributes($context);
-
-    // Check if the Entity Embed Display plugin is accessible. This also checks
-    // entity access, which is why we never call $entity->access() here.
-    if (!$display->access()) {
-      return array();
-    }
-
-    return $display->build();
-  }
-
-  /**
-   * Returns the entity manager.
-   *
-   * @return \Drupal\Core\Entity\EntityManagerInterface
-   *   The entity manager.
-   */
-  protected function entityManager() {
-    if (!isset($this->entityManager)) {
-      $this->entityManager = \Drupal::entityManager();
-    }
-    return $this->entityManager;
-  }
-
-  /**
-   * Sets the entity manager service.
-   *
-   * @param \Drupal\Core\Entity\EntityManagerInterface $entity_manager
-   *   The entity manager service.
-   *
-   * @return self
-   */
-  public function setEntityManager(EntityManagerInterface $entity_manager) {
-    $this->entityManager = $entity_manager;
-    return $this;
-  }
-
-  /**
-   * Returns the module handler.
-   *
-   * @return \Drupal\Core\Extension\ModuleHandlerInterface
-   *   The module handler.
-   */
-  protected function moduleHandler() {
-    if (!isset($this->moduleHandler)) {
-      $this->moduleHandler = \Drupal::moduleHandler();
-    }
-    return $this->moduleHandler;
-  }
-
-  /**
-   * Sets the module handler service.
-   *
-   * @param \Drupal\Core\Extension\ModuleHandlerInterface $module_handler
-   *   The module handler service.
-   *
-   * @return self
-   */
-  public function setModuleHandler(ModuleHandlerInterface $module_handler) {
-    $this->moduleHandler = $module_handler;
-    return $this;
-  }
-
-  /**
-   * Returns the Entity Embed Display plugin manager.
-   *
-   * @return \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager
-   *   The Entity Embed Display plugin manager.
-   */
-  protected function displayPluginManager() {
-    if (!isset($this->displayPluginManager)) {
-      $this->displayPluginManager = \Drupal::service('plugin.manager.entity_embed.display');
-    }
-    return $this->displayPluginManager;
-  }
-
-  /**
-   * Sets the Entity Embed Display plugin manager service.
-   *
-   * @param \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager $display_plugin_manager
-   *   The Entity Embed Display plugin manager service.
-   *
-   * @return self
-   */
-  public function setDisplayPluginManager(EntityEmbedDisplayManager $display_plugin_manager) {
-    $this->displayPluginManager = $display_plugin_manager;
-    return $this;
-  }
-
-}
diff --git a/src/Form/EntityEmbedDialog.php b/src/Form/EntityEmbedDialog.php
index 6ca9e79..ff91751 100644
--- a/src/Form/EntityEmbedDialog.php
+++ b/src/Form/EntityEmbedDialog.php
@@ -1,10 +1,5 @@
 <?php
 
-/**
- * @file
- * Contains \Drupal\entity_embed\Form\EntityEmbedDialog.
- */
-
 namespace Drupal\entity_embed\Form;
 
 use Drupal\Component\Utility\Html;
@@ -13,8 +8,10 @@ use Drupal\Core\Ajax\AjaxResponse;
 use Drupal\Core\Ajax\CloseModalDialogCommand;
 use Drupal\Core\Ajax\HtmlCommand;
 use Drupal\Core\Ajax\SetDialogTitleCommand;
+use Drupal\Core\Entity\EntityFieldManagerInterface;
 use Drupal\Core\Entity\EntityInterface;
 use Drupal\Core\Entity\EntityTypeManagerInterface;
+use Drupal\Core\Extension\ModuleHandlerInterface;
 use Drupal\Core\Form\FormBase;
 use Drupal\Core\Form\FormBuilderInterface;
 use Drupal\Core\Form\FormStateInterface;
@@ -24,7 +21,6 @@ use Drupal\embed\EmbedButtonInterface;
 use Drupal\entity_browser\Events\Events;
 use Drupal\entity_browser\Events\RegisterJSCallbacks;
 use Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager;
-use Drupal\entity_embed\EntityHelperTrait;
 use Drupal\Component\Serialization\Json;
 use Symfony\Component\DependencyInjection\ContainerInterface;
 use Symfony\Component\EventDispatcher\EventDispatcherInterface;
@@ -33,7 +29,13 @@ use Symfony\Component\EventDispatcher\EventDispatcherInterface;
  * Provides a form to embed entities by specifying data attributes.
  */
 class EntityEmbedDialog extends FormBase {
-  use EntityHelperTrait;
+
+  /**
+   * The entity embed display manager.
+   *
+   * @var \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager
+   */
+  protected $entityEmbedDisplayManager;
 
   /**
    * The form builder.
@@ -64,6 +66,20 @@ class EntityEmbedDialog extends FormBase {
   protected $entityBrowser;
 
   /**
+   * The entity field manager.
+   *
+   * @var \Drupal\Core\Entity\EntityFieldManager
+   */
+  protected $entityFieldManager;
+
+  /**
+   * The module handler service.
+   *
+   * @var \Drupal\Core\Extension\ModuleHandlerInterface
+   */
+  protected $moduleHandler;
+
+  /**
    * The entity browser settings from the entity embed button.
    */
   protected $entityBrowserSettings = [];
@@ -71,7 +87,7 @@ class EntityEmbedDialog extends FormBase {
   /**
    * Constructs a EntityEmbedDialog object.
    *
-   * @param \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager $plugin_manager
+   * @param \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager $entity_embed_display_manager
    *   The Module Handler.
    * @param \Drupal\Core\Form\FormBuilderInterface $form_builder
    *   The Form Builder.
@@ -79,12 +95,18 @@ class EntityEmbedDialog extends FormBase {
    *   The entity type manager service.
    * @param \Symfony\Component\EventDispatcher\EventDispatcherInterface $event_dispatcher
    *   Event dispatcher service.
+   * @param \Drupal\Core\Entity\EntityFieldManagerInterface $entity_field_manager
+   *   The entity field manager.
+   * @param \Drupal\Core\Extension\ModuleHandlerInterface $module_handler
+   *   The module handler.
    */
-  public function __construct(EntityEmbedDisplayManager $plugin_manager, FormBuilderInterface $form_builder, EntityTypeManagerInterface $entity_type_manager, EventDispatcherInterface $event_dispatcher) {
-    $this->setDisplayPluginManager($plugin_manager);
+  public function __construct(EntityEmbedDisplayManager $entity_embed_display_manager, FormBuilderInterface $form_builder, EntityTypeManagerInterface $entity_type_manager, EventDispatcherInterface $event_dispatcher, EntityFieldManagerInterface $entity_field_manager, ModuleHandlerInterface $module_handler) {
+    $this->entityEmbedDisplayManager = $entity_embed_display_manager;
     $this->formBuilder = $form_builder;
     $this->entityTypeManager = $entity_type_manager;
     $this->eventDispatcher = $event_dispatcher;
+    $this->entityFieldManager = $entity_field_manager;
+    $this->moduleHandler = $module_handler;
   }
 
   /**
@@ -95,7 +117,9 @@ class EntityEmbedDialog extends FormBase {
       $container->get('plugin.manager.entity_embed.display'),
       $container->get('form_builder'),
       $container->get('entity_type.manager'),
-      $container->get('event_dispatcher')
+      $container->get('event_dispatcher'),
+      $container->get('entity_field.manager'),
+      $container->get('module_handler')
     );
   }
 
@@ -137,7 +161,9 @@ class EntityEmbedDialog extends FormBase {
       'data-entity-embed-settings' => array(),
     );
     $form_state->set('entity_element', $entity_element);
-    $form_state->set('entity', $this->loadEntity($entity_element['data-entity-type'], $entity_element['data-entity-uuid']));
+    $entity = $this->entityTypeManager->getStorage($entity_element['data-entity-type'])
+      ->loadByProperties(['uuid' => $entity_element['data-entity-uuid']]);
+    $form_state->set('entity', current($entity) ?: NULL);
 
     if (!$form_state->get('step')) {
       // If an entity has been selected, then always skip to the embed options.
@@ -199,7 +225,7 @@ class EntityEmbedDialog extends FormBase {
     // the label field definition.
     $entity_type = $this->entityTypeManager->getDefinition($entity_element['data-entity-type']);
     if ($entity_type->isSubclassOf('\Drupal\Core\Entity\FieldableEntityInterface') && $entity_type->hasKey('label')) {
-      $field_definitions = $this->entityManager()->getBaseFieldDefinitions($entity_type->id());
+      $field_definitions = $this->entityFieldManager->getBaseFieldDefinitions($entity_type->id());
       if (isset($field_definitions[$entity_type->getKey('label')])) {
         $label = $field_definitions[$entity_type->getKey('label')]->getLabel();
       }
@@ -414,7 +440,7 @@ class EntityEmbedDialog extends FormBase {
       if (is_string($entity_element['data-entity-embed-settings'])) {
         $entity_element['data-entity-embed-settings'] = Json::decode($entity_element['data-entity-embed-settings']);
       }
-      $display = $this->displayPluginManager()->createInstance($plugin_id, $entity_element['data-entity-embed-settings']);
+      $display = $this->entityEmbedDisplayManager->createInstance($plugin_id, $entity_element['data-entity-embed-settings']);
       $display->setContextValue('entity', $entity);
       $display->setAttributes($entity_element);
       $form['attributes']['data-entity-embed-settings'] += $display->buildConfigurationForm($form, $form_state);
@@ -510,7 +536,8 @@ class EntityEmbedDialog extends FormBase {
     }
 
     $entity_type = $form_state->getValue(['attributes', 'data-entity-type']);
-    if ($entity = $this->loadEntity($entity_type, $id)) {
+
+    if ($entity = $this->entityTypeManager->getStorage($entity_type)->load($id)) {
       if (!$entity->access('view')) {
         $form_state->setError($element, $this->t('Unable to access @type entity @id.', array('@type' => $entity_type, '@id' => $id)));
       }
@@ -550,10 +577,12 @@ class EntityEmbedDialog extends FormBase {
   public function validateEmbedStep(array $form, FormStateInterface $form_state) {
     // Validate configuration forms for the Entity Embed Display plugin used.
     $entity_element = $form_state->getValue('attributes');
-    $entity = $this->loadEntity($entity_element['data-entity-type'], $entity_element['data-entity-uuid']);
+    $entity = $this->entityTypeManager->getStorage($entity_element['data-entity-type'])
+      ->loadByProperties(['uuid' => $entity_element['data-entity-uuid']]);
+    $entity = current($entity) ?: NULL;
     $plugin_id = $entity_element['data-entity-embed-display'];
     $plugin_settings = $entity_element['data-entity-embed-settings'] ?: array();
-    $display = $this->displayPluginManager()->createInstance($plugin_id, $plugin_settings);
+    $display = $this->entityEmbedDisplayManager->createInstance($plugin_id, $plugin_settings);
     $display->setContextValue('entity', $entity);
     $display->setAttributes($entity_element);
     $display->validateConfigurationForm($form, $form_state);
@@ -705,10 +734,12 @@ class EntityEmbedDialog extends FormBase {
 
     // Submit configuration form the selected Entity Embed Display plugin.
     $entity_element = $form_state->getValue('attributes');
-    $entity = $this->loadEntity($entity_element['data-entity-type'], $entity_element['data-entity-uuid']);
+    $entity = $this->entityTypeManager->getStorage($entity_element['data-entity-type'])
+      ->loadByProperties(['uuid' => $entity_element['data-entity-uuid']]);
+    $entity = current($entity);
     $plugin_id = $entity_element['data-entity-embed-display'];
     $plugin_settings = $entity_element['data-entity-embed-settings'] ?: array();
-    $display = $this->displayPluginManager()->createInstance($plugin_id, $plugin_settings);
+    $display = $this->entityEmbedDisplayManager->createInstance($plugin_id, $plugin_settings);
     $display->setContextValue('entity', $entity);
     $display->setAttributes($entity_element);
     $display->submitConfigurationForm($form, $form_state);
@@ -735,7 +766,7 @@ class EntityEmbedDialog extends FormBase {
       });
 
       // Allow other modules to alter the values before getting submitted to the WYSIWYG.
-      $this->moduleHandler()->alter('entity_embed_values', $values, $entity, $display, $form_state);
+      $this->moduleHandler->alter('entity_embed_values', $values, $entity, $display, $form_state);
 
       $response->addCommand(new EditorDialogSave($values));
       $response->addCommand(new CloseModalDialogCommand());
@@ -769,7 +800,7 @@ class EntityEmbedDialog extends FormBase {
    *   List of allowed Entity Embed Display plugins.
    */
   public function getDisplayPluginOptions(EmbedButtonInterface $embed_button, EntityInterface $entity) {
-    $plugins = $this->displayPluginManager->getDefinitionOptionsForEntity($entity);
+    $plugins = $this->entityEmbedDisplayManager->getDefinitionOptionsForEntity($entity);
 
     if ($allowed_plugins = $embed_button->getTypeSetting('display_plugins')) {
       $plugins = array_intersect_key($plugins, array_flip($allowed_plugins));
diff --git a/src/Plugin/Filter/EntityEmbedFilter.php b/src/Plugin/Filter/EntityEmbedFilter.php
index 965e5c9..755d746 100644
--- a/src/Plugin/Filter/EntityEmbedFilter.php
+++ b/src/Plugin/Filter/EntityEmbedFilter.php
@@ -1,20 +1,16 @@
 <?php
 
-/**
- * @file
- * Contains \Drupal\entity_embed\Plugin\Filter\EntityEmbedFilter.
- */
-
 namespace Drupal\entity_embed\Plugin\Filter;
 
 use Drupal\Component\Utility\Html;
-use Drupal\Core\Entity\EntityManagerInterface;
+use Drupal\Core\Entity\EntityTypeManagerInterface;
 use Drupal\Core\Extension\ModuleHandlerInterface;
 use Drupal\Core\Plugin\ContainerFactoryPluginInterface;
 use Drupal\Core\Render\BubbleableMetadata;
 use Drupal\Core\Render\RenderContext;
 use Drupal\Core\Render\RendererInterface;
-use Drupal\entity_embed\EntityHelperTrait;
+use Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager;
+use Drupal\entity_embed\EntityEmbedHelperTrait;
 use Drupal\entity_embed\Exception\EntityNotFoundException;
 use Drupal\entity_embed\Exception\RecursiveRenderingException;
 use Drupal\filter\FilterProcessResult;
@@ -33,7 +29,8 @@ use Drupal\embed\DomHelperTrait;
  * )
  */
 class EntityEmbedFilter extends FilterBase implements ContainerFactoryPluginInterface {
-  use EntityHelperTrait;
+
+  use EntityEmbedHelperTrait;
   use DomHelperTrait;
 
   /**
@@ -44,6 +41,27 @@ class EntityEmbedFilter extends FilterBase implements ContainerFactoryPluginInte
   protected $renderer;
 
   /**
+   * The entity type manager.
+   *
+   * @var \Drupal\Core\Entity\EntityTypeManagerInterface
+   */
+  protected $entityTypeManager;
+
+  /**
+   * The module handler service.
+   *
+   * @var \Drupal\Core\Extension\ModuleHandlerInterface
+   */
+  protected $moduleHandler;
+
+  /**
+   * The entity embed display plugin manager service.
+   *
+   * @var \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager
+   */
+  protected $displayPluginManager;
+
+  /**
    * Constructs a EntityEmbedFilter object.
    *
    * @param array $configuration
@@ -52,16 +70,17 @@ class EntityEmbedFilter extends FilterBase implements ContainerFactoryPluginInte
    *   The plugin ID for the plugin instance.
    * @param mixed $plugin_definition
    *   The plugin implementation definition.
-   * @param \Drupal\Core\Entity\EntityManagerInterface $entity_manager
-   *   The entity manager service.
-   * @param \Drupal\Core\Extension\ModuleHandlerInterface $module_handler
-   *   The Module Handler.
+   * @param \Drupal\Core\Entity\EntityTypeManagerInterface $entity_type_manager
+   *   The entity type manager.
+   * @param \Drupal\Core\Render\RendererInterface $renderer
+   *   The renderer.
    */
-  public function __construct(array $configuration, $plugin_id, $plugin_definition, EntityManagerInterface $entity_manager, ModuleHandlerInterface $module_handler, RendererInterface $renderer) {
+  public function __construct(array $configuration, $plugin_id, $plugin_definition, EntityTypeManagerInterface $entity_type_manager, RendererInterface $renderer, ModuleHandlerInterface $module_handler, EntityEmbedDisplayManager $display_manager) {
     parent::__construct($configuration, $plugin_id, $plugin_definition);
-    $this->setEntityManager($entity_manager);
-    $this->setModuleHandler($module_handler);
+    $this->entityTypeManager = $entity_type_manager;
     $this->renderer = $renderer;
+    $this->moduleHandler = $module_handler;
+    $this->displayPluginManager = $display_manager;
   }
 
   /**
@@ -72,9 +91,10 @@ class EntityEmbedFilter extends FilterBase implements ContainerFactoryPluginInte
       $configuration,
       $plugin_id,
       $plugin_definition,
-      $container->get('entity.manager'),
+      $container->get('entity_type.manager'),
+      $container->get('renderer'),
       $container->get('module_handler'),
-      $container->get('renderer')
+      $container->get('plugin.manager.entity_embed.display')
     );
   }
 
@@ -96,8 +116,17 @@ class EntityEmbedFilter extends FilterBase implements ContainerFactoryPluginInte
 
         try {
           // Load the entity either by UUID (preferred) or ID.
-          $id = $node->getAttribute('data-entity-uuid') ?: $node->getAttribute('data-entity-id');
-          $entity = $this->loadEntity($entity_type, $id);
+          $id = NULL;
+          $entity = NULL;
+          if ($id = $node->getAttribute('data-entity-uuid')) {
+            $entity = $this->entityTypeManager->getStorage($entity_type)
+              ->loadByProperties(['uuid' => $id]);
+            $entity = current($entity);
+          }
+          else {
+            $id = $node->getAttribute('data-entity-id');
+            $entity = $this->entityTypeManager->getStorage($entity_type)->load($id);
+          }
 
           if ($entity) {
             // Protect ourselves from recursive rendering.
diff --git a/src/Plugin/entity_embed/EntityEmbedDisplay/ImageFieldFormatter.php b/src/Plugin/entity_embed/EntityEmbedDisplay/ImageFieldFormatter.php
index c1f2736..ec2d60e 100644
--- a/src/Plugin/entity_embed/EntityEmbedDisplay/ImageFieldFormatter.php
+++ b/src/Plugin/entity_embed/EntityEmbedDisplay/ImageFieldFormatter.php
@@ -1,17 +1,13 @@
 <?php
 
-/**
- * @file
- * Contains \Drupal\entity_embed\Plugin\entity_embed\EntityEmbedDisplay\ImageFieldFormatter.
- */
-
 namespace Drupal\entity_embed\Plugin\entity_embed\EntityEmbedDisplay;
 
 use Drupal\Core\Access\AccessResult;
-use Drupal\Core\Entity\EntityManagerInterface;
+use Drupal\Core\Entity\EntityTypeManagerInterface;
 use Drupal\Core\Field\FormatterPluginManager;
 use Drupal\Core\Form\FormStateInterface;
 use Drupal\Core\Image\ImageFactory;
+use Drupal\Core\Language\LanguageManagerInterface;
 use Drupal\Core\Session\AccountInterface;
 use Drupal\Core\TypedData\TypedDataManager;
 use Symfony\Component\DependencyInjection\ContainerInterface;
@@ -33,16 +29,16 @@ use Symfony\Component\DependencyInjection\ContainerInterface;
 class ImageFieldFormatter extends FileFieldFormatter {
 
   /**
-    * The image factory.
-    *
-    * @var \Drupal\Core\Image\ImageFactory
-    */
-    protected $imageFactory;
+   * The image factory.
+   *
+   * @var \Drupal\Core\Image\ImageFactory
+   */
+  protected $imageFactory;
 
   /**
    * {@inheritdoc}
    *
-   * @param \Drupal\Core\Entity\EntityManagerInterface $entity_manager
+   * @param \Drupal\Core\Entity\EntityTypeManagerInterface $entity_type_manager
    *   The entity manager service.
    * @param \Drupal\Core\Field\FormatterPluginManager $formatter_plugin_manager
    *   The field formatter plugin manager.
@@ -50,9 +46,11 @@ class ImageFieldFormatter extends FileFieldFormatter {
    *   The typed data manager.
    * @param \Drupal\Core\Image\ImageFactory $image_factory
    *   The image factory.
+   * @param \Drupal\Core\Language\LanguageManagerInterface $language_manager
+   *   The language manager.
    */
-  public function __construct(array $configuration, $plugin_id, $plugin_definition, EntityManagerInterface $entity_manager, FormatterPluginManager $formatter_plugin_manager, TypedDataManager $typed_data_manager, ImageFactory $image_factory) {
-    parent::__construct($configuration, $plugin_id, $plugin_definition, $entity_manager, $formatter_plugin_manager, $typed_data_manager);
+  public function __construct(array $configuration, $plugin_id, $plugin_definition, EntityTypeManagerInterface $entity_type_manager, FormatterPluginManager $formatter_plugin_manager, TypedDataManager $typed_data_manager, ImageFactory $image_factory, LanguageManagerInterface $language_manager) {
+    parent::__construct($configuration, $plugin_id, $plugin_definition, $entity_type_manager, $formatter_plugin_manager, $typed_data_manager, $language_manager);
     $this->imageFactory = $image_factory;
   }
 
@@ -64,10 +62,11 @@ class ImageFieldFormatter extends FileFieldFormatter {
       $configuration,
       $plugin_id,
       $plugin_definition,
-      $container->get('entity.manager'),
+      $container->get('entity_type.manager'),
       $container->get('plugin.manager.field.formatter'),
       $container->get('typed_data_manager'),
-      $container->get('image.factory')
+      $container->get('image.factory'),
+      $container->get('language_manager')
     );
   }
 
diff --git a/src/Tests/EntityEmbedDialogTest.php b/src/Tests/EntityEmbedDialogTest.php
index bbcdc0f..62f1474 100644
--- a/src/Tests/EntityEmbedDialogTest.php
+++ b/src/Tests/EntityEmbedDialogTest.php
@@ -110,7 +110,7 @@ class EntityEmbedDialogTest extends EntityEmbedTestBase {
 
     $this->container->get('config.factory')->getEditable('entity_embed.settings')
       ->set('rendered_entity_mode', TRUE)->save();
-    $this->displayPluginManager()->clearCachedDefinitions();
+    $this->container->get('plugin.manager.entity_embed.display')->clearCachedDefinitions();
 
     $this->getEmbedDialog('custom_format', 'node');
     $title = $this->node->getTitle() . ' (' . $this->node->id() . ')';
diff --git a/src/Tests/EntityEmbedHooksTest.php b/src/Tests/EntityEmbedHooksTest.php
index 714a89f..9ed096d 100644
--- a/src/Tests/EntityEmbedHooksTest.php
+++ b/src/Tests/EntityEmbedHooksTest.php
@@ -1,10 +1,5 @@
 <?php
 
-/**
- * @file
- * Contains \Drupal\entity_embed\Tests\EntityEmbedHooksTest.
- */
-
 namespace Drupal\entity_embed\Tests;
 
 /**
@@ -34,7 +29,8 @@ class EntityEmbedHooksTest extends EntityEmbedTestBase {
     // hook_entity_embed_display_plugins_alter() implementation and ensure it is
     // working as designed.
     $this->state->set('entity_embed_test_entity_embed_display_plugins_alter', TRUE);
-    $plugins = $this->displayPluginManager()->getDefinitionOptionsForEntity($this->node);
+    $plugins = $this->container->get('plugin.manager.entity_embed.display')
+      ->getDefinitionOptionsForEntity($this->node);
     // Ensure that name of each plugin is prefixed with 'testing_hook:'.
     foreach ($plugins as $plugin => $plugin_info) {
       $this->assertTrue(strpos($plugin, 'testing_hook:') === 0, 'Name of the plugin is prefixed by hook_entity_embed_display_plugins_alter()');
diff --git a/src/Tests/EntityEmbedTestBase.php b/src/Tests/EntityEmbedTestBase.php
index 0958a67..b4c93ca 100644
--- a/src/Tests/EntityEmbedTestBase.php
+++ b/src/Tests/EntityEmbedTestBase.php
@@ -1,15 +1,9 @@
 <?php
 
-/**
- * @file
- * Contains \Drupal\entity_embed\Tests\EntityEmbedTestBase.
- */
-
 namespace Drupal\entity_embed\Tests;
 
 use Drupal\Core\Entity\EntityInterface;
 use Drupal\editor\Entity\Editor;
-use Drupal\entity_embed\EntityHelperTrait;
 use Drupal\file\Entity\File;
 use Drupal\filter\Entity\FilterFormat;
 use Drupal\simpletest\WebTestBase;
@@ -18,7 +12,6 @@ use Drupal\simpletest\WebTestBase;
  * Base class for all entity_embed tests.
  */
 abstract class EntityEmbedTestBase extends WebTestBase {
-  use EntityHelperTrait;
 
   /**
    * Modules to enable.
@@ -111,10 +104,12 @@ abstract class EntityEmbedTestBase extends WebTestBase {
   }
 
   public function assertAvailableDisplayPlugins(EntityInterface $entity, array $expected_plugins, $message = '') {
-    $plugin_options = $this->displayPluginManager()->getDefinitionOptionsForEntity($entity);
+    $plugin_options = $this->container->get('plugin.manager.entity_embed.display')
+      ->getDefinitionOptionsForEntity($entity);
     // @todo Remove the sorting so we can actually test return order.
     ksort($plugin_options);
     sort($expected_plugins);
     $this->assertEqual(array_keys($plugin_options), $expected_plugins, $message);
   }
+
 }
diff --git a/src/Tests/EntityReferenceFieldFormatterTest.php b/src/Tests/EntityReferenceFieldFormatterTest.php
index 9dbe32a..6d8e770 100644
--- a/src/Tests/EntityReferenceFieldFormatterTest.php
+++ b/src/Tests/EntityReferenceFieldFormatterTest.php
@@ -47,7 +47,7 @@ class EntityReferenceFieldFormatterTest extends EntityEmbedTestBase {
 
     $this->container->get('config.factory')->getEditable('entity_embed.settings')
       ->set('rendered_entity_mode', TRUE)->save();
-    $this->displayPluginManager()->clearCachedDefinitions();
+    $this->container->get('plugin.manager.entity_embed.display')->clearCachedDefinitions();
 
     $this->assertAvailableDisplayPlugins($this->node, [
       'entity_reference:entity_reference_label',
@@ -59,7 +59,7 @@ class EntityReferenceFieldFormatterTest extends EntityEmbedTestBase {
     // 'entity_reference:entity_reference_entity_id' plugin.
     $form = array();
     $form_state = new FormState();
-    $display = $this->displayPluginManager()->createInstance('entity_reference:entity_reference_entity_id', array());
+    $display = $this->container->get('plugin.manager.entity_embed.display')->createInstance('entity_reference:entity_reference_entity_id', array());
     $display->setContextValue('entity', $this->node);
     $conf_form = $display->buildConfigurationForm($form, $form_state);
     $this->assertIdentical(array_keys($conf_form), array());
@@ -68,7 +68,7 @@ class EntityReferenceFieldFormatterTest extends EntityEmbedTestBase {
     // 'entity_reference:entity_reference_entity_view' plugin.
     $form = array();
     $form_state = new FormState();
-    $display = $this->displayPluginManager()->createInstance('entity_reference:entity_reference_entity_view', array());
+    $display = $this->container->get('plugin.manager.entity_embed.display')->createInstance('entity_reference:entity_reference_entity_view', array());
     $display->setContextValue('entity', $this->node);
     $conf_form = $display->buildConfigurationForm($form, $form_state);
     $this->assertIdentical($conf_form['view_mode']['#type'], 'select');
@@ -78,7 +78,7 @@ class EntityReferenceFieldFormatterTest extends EntityEmbedTestBase {
     // 'entity_reference:entity_reference_label' plugin.
     $form = array();
     $form_state = new FormState();
-    $display = $this->displayPluginManager()->createInstance('entity_reference:entity_reference_label', array());
+    $display = $this->container->get('plugin.manager.entity_embed.display')->createInstance('entity_reference:entity_reference_label', array());
     $display->setContextValue('entity', $this->node);
     $conf_form = $display->buildConfigurationForm($form, $form_state);
     $this->assertIdentical(array_keys($conf_form), array('link'));
@@ -87,7 +87,7 @@ class EntityReferenceFieldFormatterTest extends EntityEmbedTestBase {
 
     // Ensure that 'Rendered Entity' plugin is not available for an entity not
     // having a view controller.
-    $plugin_options = $this->displayPluginManager()->getDefinitionOptionsForEntity($this->menu);
+    $plugin_options = $this->container->get('plugin.manager.entity_embed.display')->getDefinitionOptionsForEntity($this->menu);
     $this->assertFalse(array_key_exists('entity_reference:entity_reference_entity_view', $plugin_options), "The 'Rendered entity' plugin is not available.");
   }
 
diff --git a/src/Tests/FileFieldFormatterTest.php b/src/Tests/FileFieldFormatterTest.php
index 01b6118..45e7573 100644
--- a/src/Tests/FileFieldFormatterTest.php
+++ b/src/Tests/FileFieldFormatterTest.php
@@ -1,10 +1,5 @@
 <?php
 
-/**
- * @file
- * Contains \Drupal\entity_embed\Tests\FileFieldFormatterTest.
- */
-
 namespace Drupal\entity_embed\Tests;
 
 use Drupal\Component\Serialization\Json;
@@ -60,7 +55,8 @@ class FileFieldFormatterTest extends EntityEmbedTestBase {
     );
     // Ensure that description field is available for all the 'file' plugins.
     foreach ($plugins as $plugin) {
-      $display = $this->displayPluginManager()->createInstance($plugin, array());
+      $display = $this->container->get('plugin.manager.entity_embed.display')
+        ->createInstance($plugin, []);
       $display->setContextValue('entity', $this->file);
       $conf_form = $display->buildConfigurationForm($form, $form_state);
       $this->assertIdentical(array_keys($conf_form), array('description'));
diff --git a/src/Tests/ImageFieldFormatterTest.php b/src/Tests/ImageFieldFormatterTest.php
index de2f6ff..96b86e7 100644
--- a/src/Tests/ImageFieldFormatterTest.php
+++ b/src/Tests/ImageFieldFormatterTest.php
@@ -1,10 +1,5 @@
 <?php
 
-/**
- * @file
- * Contains \Drupal\entity_embed\Tests\ImageFieldFormatterTest.
- */
-
 namespace Drupal\entity_embed\Tests;
 
 use Drupal\Component\Serialization\Json;
@@ -62,7 +57,8 @@ class ImageFieldFormatterTest extends EntityEmbedTestBase {
     // Ensure that correct form attributes are returned for the image plugin.
     $form = array();
     $form_state = new FormState();
-    $display = $this->displayPluginManager()->createInstance('image:image', array());
+    $display = $this->container->get('plugin.manager.entity_embed.display')
+      ->createInstance('image:image', []);
     $display->setContextValue('entity', $this->image);
     $conf_form = $display->buildConfigurationForm($form, $form_state);
     $this->assertIdentical(array_keys($conf_form), array(
diff --git a/src/Tests/ViewModeFieldFormatterTest.php b/src/Tests/ViewModeFieldFormatterTest.php
index 0891901..7712da9 100644
--- a/src/Tests/ViewModeFieldFormatterTest.php
+++ b/src/Tests/ViewModeFieldFormatterTest.php
@@ -27,7 +27,8 @@ class ViewModeFieldFormatterTest extends EntityEmbedTestBase {
     foreach ($this->plugins as $plugin) {
       $form = [];
       $form_state = new FormState();
-      $display = $this->displayPluginManager()->createInstance($plugin, []);
+      $display = $this->container->get('plugin.manager.entity_embed.display')
+        ->createInstance($plugin, []);
       $display->setContextValue('entity', $this->node);
       $conf_form = $display->buildConfigurationForm($form, $form_state);
       $this->assertIdentical(array_keys($conf_form), []);
diff --git a/src/Twig/EntityEmbedTwigExtension.php b/src/Twig/EntityEmbedTwigExtension.php
index 518088c..d188bc6 100644
--- a/src/Twig/EntityEmbedTwigExtension.php
+++ b/src/Twig/EntityEmbedTwigExtension.php
@@ -1,37 +1,63 @@
 <?php
 
-/**
- * @file
- * Contains \Drupal\entity_embed\Twig\EntityEmbedTwigExtension.
- */
-
 namespace Drupal\entity_embed\Twig;
 
-use Drupal\Core\Entity\EntityManagerInterface;
-use Drupal\entity_embed\EntityHelperTrait;
+use Drupal\Core\DependencyInjection\ContainerInjectionInterface;
+use Drupal\Core\Entity\EntityTypeManagerInterface;
 use Drupal\Core\Extension\ModuleHandlerInterface;
 use Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager;
+use Drupal\entity_embed\EntityEmbedHelperTrait;
+use Symfony\Component\DependencyInjection\ContainerInterface;
 
 /**
  * Provide entity embedding function within Twig templates.
  */
-class EntityEmbedTwigExtension extends \Twig_Extension {
-  use EntityHelperTrait;
+class EntityEmbedTwigExtension extends \Twig_Extension implements ContainerInjectionInterface {
+
+  use EntityEmbedHelperTrait;
+
+  /**
+   * The entity type manager service.
+   *
+   * @var \Drupal\Core\Entity\EntityTypeManagerInterface
+   */
+  protected $entityTypeManager;
+
+  /**
+   * The module handler service.
+   *
+   * @var \Drupal\Core\Extension\ModuleHandlerInterface
+   */
+  protected $moduleHandler;
+
+  /**
+   * The entity embed display plugin manager service.
+   *
+   * @var \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager
+   */
+  protected $displayPluginManager;
 
   /**
    * Constructs a new EntityEmbedTwigExtension.
    *
-   * @param \Drupal\Core\Entity\EntityManagerInterface $entity_manager
-   *   The entity manager service.
-   * @param \Drupal\Core\Extension\ModuleHandlerInterface $module_handler
-   *   The module handler.
-   * @param \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager $plugin_manager
-   *   The plugin manager.
+   * @param \Drupal\Core\Entity\EntityTypeManagerInterface $entity_type_manager
+   *   The entity type manager service.
    */
-  public function __construct(EntityManagerInterface $entity_manager, ModuleHandlerInterface $module_handler, EntityEmbedDisplayManager $plugin_manager) {
-    $this->setEntityManager($entity_manager);
-    $this->setModuleHandler($module_handler);
-    $this->setDisplayPluginManager($plugin_manager);
+  public function __construct(EntityTypeManagerInterface $entity_type_manager, ModuleHandlerInterface $module_handler, EntityEmbedDisplayManager $display_manager) {
+    $this->entityTypeManager = $entity_type_manager;
+    $this->moduleHandler = $module_handler;
+    $this->displayPluginManager = $display_manager;
+  }
+
+  /**
+   * {@inheritdoc}
+   */
+  public static function create(ContainerInterface $container) {
+    return new static(
+      $container->get('entity_type.manager'),
+      $container->get('module_handler'),
+      $container->get('plugin.manager.entity_embed.display')
+    );
   }
 
   /**
@@ -56,7 +82,7 @@ class EntityEmbedTwigExtension extends \Twig_Extension {
    * @param string $entity_type
    *   The machine name of an entity_type like 'node'.
    * @param string $entity_id
-   *   The entity ID or entity UUID.
+   *   The entity ID.
    * @param string $display_plugin
    *   (optional) The Entity Embed Display plugin to be used to render the
    *   entity.
@@ -67,7 +93,7 @@ class EntityEmbedTwigExtension extends \Twig_Extension {
    *   A render array from entity_view().
    */
   public function getRenderArray($entity_type, $entity_id, $display_plugin = 'default', array $display_settings = []) {
-    $entity = $this->loadEntity($entity_type, $entity_id);
+    $entity = $this->entityTypeManager->getStorage($entity_type)->load($entity_id);
     $context = array(
       'data-entity-type' => $entity_type,
       'data-entity-uuid' => $entity->uuid(),
