diff --git a/src/EntityEmbedDisplayBase.php b/src/EntityEmbedDisplayBase.php
index c15a1a61..037f7807 100644
--- a/src/EntityEmbedDisplayBase.php
+++ b/src/EntityEmbedDisplayBase.php
@@ -9,7 +9,7 @@
 
 use Drupal\Component\Utility\NestedArray;
 use Drupal\Core\Entity\EntityInterface;
-use Drupal\Core\Plugin\PluginBase;
+use Drupal\Core\Plugin\ContextAwarePluginBase;
 use Drupal\Core\Session\AccountInterface;
 
 /**
@@ -17,21 +17,14 @@
  *
  * @ingroup entity_embed_api
  */
-abstract class EntityEmbedDisplayBase extends PluginBase implements EntityEmbedDisplayInterface {
+abstract class EntityEmbedDisplayBase extends ContextAwarePluginBase implements EntityEmbedDisplayInterface {
 
-  /**
-   * The entity being embedded.
-   *
-   * @var \Drupal\Core\Entity\EntityInterface
-   */
-  public $entity;
-
-  /**
-   * The context for the embedded entity.
-   *
-   * @var array
-   */
-  public $context = array();
+ /**
+  * The attributes on the embedded entity.
+  *
+  * @var array
+  */
+  public $attributes = array();
 
   /**
    * {@inheritdoc}
@@ -41,32 +34,22 @@ public function __construct(array $configuration, $plugin_id, $plugin_definition
     $this->setConfiguration($configuration);
   }
 
-  public function setEntity(EntityInterface $entity) {
-    $this->entity = $entity;
-    return $this;
-  }
-
-  public function getEntity() {
-    return $this->entity;
-  }
-
-  public function setContext(array $context) {
-    $this->context = $context;
-    return $this;
+  public function getConfigurationValue($name, $default = NULL) {
+    $configuration = $this->getConfiguration();
+    return array_key_exists($name, $configuration) ? $configuration[$name] : $default;
   }
 
-  public function getContext() {
-    return $this->context;
+  public function setAttributes(array $attributes) {
+    $this->attributes = $attributes;
   }
 
-  public function getContextValue($name, $default = NULL) {
-    $context = $this->getContext();
-    return array_key_exists($name, $context) ? $context[$name] : $default;
+  public function getAttributes() {
+    return $this->attributes;
   }
 
-  public function getConfigurationValue($name, $default = NULL) {
-    $configuration = $this->getConfiguration();
-    return array_key_exists($name, $configuration) ? $configuration[$name] : $default;
+  public function getAttributeValue($name, $default = NULL) {
+    $attributes = $this->getAttributes();
+    return array_key_exists($name, $attributes) ? $attributes[$name] : $default;
   }
 
   /**
@@ -82,7 +65,7 @@ public function access(AccountInterface $account = NULL) {
     }
 
     // Check that the entity itself can be viewed by the user.
-    return $this->entity->access('view', $account);
+    return $this->getContextValue('entity')->access('view', $account);
   }
 
   /**
diff --git a/src/EntityEmbedDisplayManager.php b/src/EntityEmbedDisplayManager.php
index 0543831e..943c29e5 100644
--- a/src/EntityEmbedDisplayManager.php
+++ b/src/EntityEmbedDisplayManager.php
@@ -14,6 +14,7 @@
 use Drupal\Core\Plugin\DefaultPluginManager;
 use Drupal\Core\Plugin\Discovery\AlterDecorator;
 use Drupal\Core\Plugin\Discovery\AnnotatedClassDiscovery;
+use Drupal\Component\Plugin\Exception\PluginException;
 
 /**
  * Provides an Entity embed display plugin manager.
@@ -69,9 +70,14 @@ public function processDefinition(&$definition, $plugin_id) {
   public function getDefinitionsByEntity(EntityInterface $entity) {
     $definitions = $this->getDefinitions();
     $valid_ids = array_filter(array_keys($definitions), function ($id) use ($entity) {
-      $display = $this->createInstance($id);
-      $display->setEntity($entity);
-      return $display->access();
+      try {
+        $display = $this->createInstance($id);
+        $display->setContextValue('entity', $entity);
+        return $display->access();
+      }
+      catch (PluginException $e) {
+        return FALSE;
+      }
     });
     return array_intersect_key($definitions, array_flip($valid_ids));
   }
diff --git a/src/FieldFormatterEntityEmbedDisplayBase.php b/src/FieldFormatterEntityEmbedDisplayBase.php
index 51a692f4..a7d3e625 100644
--- a/src/FieldFormatterEntityEmbedDisplayBase.php
+++ b/src/FieldFormatterEntityEmbedDisplayBase.php
@@ -30,7 +30,7 @@ public function getFormatter(FieldDefinition $definition = NULL) {
     }
 
     // Ensure that the field name is unique each time this is run.
-    $definition->setName('_entity_embed_' . $this->getContextValue('token'));
+    $definition->setName('_entity_embed_' . $this->getAttributeValue('token'));
 
     $display = array(
       'type' => $this->getDerivativeId(),
@@ -61,7 +61,7 @@ public function build() {
     // Field name is only set to avoid broken CSS classes.
     $definition = $this->getFieldDefinition();
     // Ensure that the field name is unique each time this is run.
-    $definition->setName('_entity_embed_' . $this->getContextValue('token'));
+    $definition->setName('_entity_embed_' . $this->getAttributeValue('token'));
 
     /* @var \Drupal\Core\Field\FieldItemListInterface $items $items */
     // Create a field item list object, 1 is the value, array('target_id' => 1)
@@ -74,7 +74,7 @@ public function build() {
       $node
     );
 
-    if ($langcode = $this->getContextValue('langcode')) {
+    if ($langcode = $this->getAttributeValue('langcode')) {
       $items->setLangcode($langcode);
     }
 
diff --git a/src/Plugin/EntityEmbedDisplay/EntityEmbedDefaultDisplay.php b/src/Plugin/EntityEmbedDisplay/EntityEmbedDefaultDisplay.php
index eaa71d1b..72f1301a 100644
--- a/src/Plugin/EntityEmbedDisplay/EntityEmbedDefaultDisplay.php
+++ b/src/Plugin/EntityEmbedDisplay/EntityEmbedDefaultDisplay.php
@@ -14,8 +14,15 @@
  *
  * @EntityEmbedDisplay(
  *   id = "default",
- *   label = @Translation("Default")
+ *   label = @Translation("Default"),
+ *   context = {
+ *     "entity" = {
+ *       "type" = "entity"
+ *     }
+ *   }
  * )
+ *
+ * @todo Should this use a derivative? http://cgit.drupalcode.org/page_manager/tree/src/Plugin/Deriver/EntityViewDeriver.php
  */
 class EntityEmbedDefaultDisplay extends EntityEmbedDisplayBase {
 
@@ -37,7 +44,7 @@ public function buildConfigurationForm(array $form, array &$form_state) {
     $form['view_mode'] = array(
       '#type' => 'select',
       '#title' => t('View mode'),
-      '#options' => \Drupal::entityManager()->getDisplayModeOptions('view_mode', $this->entity->getEntityTypeId()),
+      '#options' => \Drupal::entityManager()->getDisplayModeOptions('view_mode', $this->getAttributeValue('entity-type')),
       '#default_value' => $this->getConfigurationValue('view_mode'),
       '#required' => TRUE,
     );
@@ -52,10 +59,10 @@ public function build() {
     // Clone the entity since we're going to set some additional properties we
     // don't want kept around afterwards.
     $entity = clone $this->entity;
-    $entity->entity_embed_context = $this->getContext();
+    $entity->entity_embed_attributes = $this->getAttributes();
 
     // Build the rendered entity.
-    $build = entity_view($this->entity, $this->getConfigurationValue('view_mode'), $this->getContextValue('langcode'));
+    $build = entity_view($this->entity, $this->getConfigurationValue('view_mode'), $this->getAttributeValue('langcode'));
 
     // Hide entity links by default.
     if (isset($build['links'])) {
diff --git a/src/Plugin/EntityEmbedDisplay/EntityReferenceFieldFormatter.php b/src/Plugin/EntityEmbedDisplay/EntityReferenceFieldFormatter.php
index 5fc52ce9..debab155 100644
--- a/src/Plugin/EntityEmbedDisplay/EntityReferenceFieldFormatter.php
+++ b/src/Plugin/EntityEmbedDisplay/EntityReferenceFieldFormatter.php
@@ -16,6 +16,11 @@
  * @EntityEmbedDisplay(
  *   id = "entity_reference",
  *   label = @Translation("Entity Reference"),
+ *   context = {
+ *     "entity" = {
+ *       "type" = "entity"
+ *     }
+ *   },
  *   derivative = "Drupal\entity_embed\Plugin\Derivative\FieldFormatterDeriver",
  *   field_type = "entity_reference",
  *   provider = "entity_reference"
@@ -28,7 +33,7 @@ class EntityReferenceFieldFormatter extends FieldFormatterEntityEmbedDisplayBase
    */
   public function getFieldDefinition() {
     $field = FieldDefinition::create('entity_reference');
-    $field->setSetting('target_type', $this->getContextValue('entity-type'));
+    $field->setSetting('target_type', $this->getAttributeValue('entity-type'));
     return $field;
   }
 
@@ -36,7 +41,7 @@ public function getFieldDefinition() {
    * {@inheritdoc}
    */
   public function getFieldValue(FieldDefinition $definition) {
-    return array('target_id' => $this->entity->id());
+    return array('target_id' => $this->getContextValue('entity')->id());
   }
 
 }
diff --git a/src/Plugin/EntityEmbedDisplay/FileFieldFormatter.php b/src/Plugin/EntityEmbedDisplay/FileFieldFormatter.php
index c6df40ee..5cb068b9 100644
--- a/src/Plugin/EntityEmbedDisplay/FileFieldFormatter.php
+++ b/src/Plugin/EntityEmbedDisplay/FileFieldFormatter.php
@@ -16,6 +16,11 @@
  * @EntityEmbedDisplay(
  *   id = "file",
  *   label = @Translation("File"),
+ *   context = {
+ *     "entity" = {
+ *       "type" = "entity:file"
+ *     }
+ *   },
  *   entity_types = {"file"},
  *   derivative = "Drupal\entity_embed\Plugin\Derivative\FieldFormatterDeriver",
  *   field_type = "file",
@@ -37,7 +42,7 @@ public function getFieldDefinition() {
    */
   public function getFieldValue(FieldDefinition $definition) {
     $value = parent::getFieldValue($definition);
-    $value += array_intersect_key($this->getContext(), array('description' => ''));
+    $value += array_intersect_key($this->getAttributes(), array('description' => ''));
     return $value;
   }
 
@@ -45,6 +50,10 @@ public function getFieldValue(FieldDefinition $definition) {
    * {@inheritdoc}
    */
   public function access(AccountInterface $account = NULL) {
+<<<<<<< HEAD
+    // Due to issues with access checking with file entities in core, we must
+    // manually check hook_file_download to see if the user can access the file.
+=======
     if (!$this->isValidEntityType()) {
       return FALSE;
     }
@@ -52,10 +61,11 @@ public function access(AccountInterface $account = NULL) {
     // Due to issues with access checking with file entities in core, we cannot
     // actually use Entity::access() which would have been called by
     // parent::access().
+>>>>>>> embed-display
     // @see https://drupal.org/node/2128791
     // @see https://drupal.org/node/2148353
     // @see https://drupal.org/node/2078473
-    switch (file_uri_scheme($this->entity->getFileUri())) {
+    switch (file_uri_scheme($this->getContextValue('entity')->getFileUri())) {
       case 'public':
         return TRUE;
 
diff --git a/src/Plugin/EntityEmbedDisplay/ImageFieldFormatter.php b/src/Plugin/EntityEmbedDisplay/ImageFieldFormatter.php
index f0196ef5..e17278b3 100644
--- a/src/Plugin/EntityEmbedDisplay/ImageFieldFormatter.php
+++ b/src/Plugin/EntityEmbedDisplay/ImageFieldFormatter.php
@@ -16,6 +16,11 @@
  * @EntityEmbedDisplay(
  *   id = "image",
  *   label = @Translation("Image"),
+ *   context = {
+ *     "entity" = {
+ *       "type" = "entity:file"
+ *     }
+ *   },
  *   entity_types = {"file"},
  *   derivative = "Drupal\entity_embed\Plugin\Derivative\FieldFormatterDeriver",
  *   field_type = "image",
@@ -39,7 +44,7 @@ public function getFieldValue(FieldDefinition $definition) {
     $value = parent::getFieldValue($definition);
     // File field support descriptions, but images do not.
     unset($value['description']);
-    $value += array_intersect_key($this->getContext(), array('alt' => '', 'title' => ''));
+    $value += array_intersect_key($this->getAttributes(), array('alt' => '', 'title' => ''));
     return $value;
   }
 
@@ -51,7 +56,8 @@ public function access(AccountInterface $account = NULL) {
       return FALSE;
     }
 
-    $image = \Drupal::service('image.factory')->get($this->entity->getFileUri());
+    $uri = $this->getContextValue('entity')->getFileUri();
+    $image = \Drupal::service('image.factory')->get($uri);
     return $image->isSupported();
   }
 
@@ -74,7 +80,7 @@ public function buildConfigurationForm(array $form, array &$form_state) {
     $form['alt'] = array(
       '#type' => 'textfield',
       '#title' => $this->t('Alternate text'),
-      '#default_value' => $this->getContextValue('alt', ''),
+      '#default_value' => $this->getAttributeValue('alt', ''),
       '#description' => $this->t('This text will be used by screen readers, search engines, or when the image cannot be loaded.'),
       '#parents' => array('attributes', 'alt'),
       // @see http://www.gawds.org/show.php?contentid=28
@@ -83,7 +89,7 @@ public function buildConfigurationForm(array $form, array &$form_state) {
     $form['title'] = array(
       '#type' => 'textfield',
       '#title' => $this->t('Title'),
-      '#default_value' => $this->getContextValue('title', ''),
+      '#default_value' => $this->getAttributeValue('title', ''),
       '#description' => t('The title is used as a tool tip when the user hovers the mouse over the image.'),
       '#parents' => array('attributes', 'title'),
       '#maxlength' => 1024,
diff --git a/src/Plugin/Filter/EntityEmbedFilter.php b/src/Plugin/Filter/EntityEmbedFilter.php
index c4298c4f..c7006044 100644
--- a/src/Plugin/Filter/EntityEmbedFilter.php
+++ b/src/Plugin/Filter/EntityEmbedFilter.php
@@ -256,8 +256,8 @@ public static function postRender(array $element, array $context) {
       if ($entity = entity_load($context['entity-type'], $context['entity-id'])) {
         $manager = \Drupal::service('plugin.manager.entity_embed.display');
         $display = $manager->createInstance($context['entity-embed-display'], $context['entity-embed-settings']);
-        $display->setEntity($entity);
-        $display->setContext($context);
+        $display->setContextValue('entity', $entity);
+        $display->setAttributes($context);
         if ($display->access()) {
           $build = $display->build();
           // Allow modules to alter the rendered embedded entity.