diff --git a/src/Plugin/Block/EntityEditFormBlock.php b/src/Plugin/Block/EntityEditFormBlock.php
index 4ee3a4d..d53d513 100644
--- a/src/Plugin/Block/EntityEditFormBlock.php
+++ b/src/Plugin/Block/EntityEditFormBlock.php
@@ -3,9 +3,19 @@
 namespace Drupal\entityform_block\Plugin\Block;
 
 use Drupal\Core\Block\BlockBase;
+use Drupal\Core\Entity\Entity\EntityFormDisplay;
+use Drupal\Core\Entity\EntityDisplayRepositoryInterface;
+use Drupal\Core\Entity\EntityTypeBundleInfoInterface;
+use Drupal\Core\Entity\EntityTypeManagerInterface;
+use Drupal\Core\Entity\FieldableEntityInterface;
+use Drupal\Core\Form\FormBuilder;
+use Drupal\Core\Form\FormState;
 use Drupal\Core\Form\FormStateInterface;
+use Drupal\Core\Plugin\ContainerFactoryPluginInterface;
+use Drupal\Core\Plugin\ContextAwarePluginInterface;
 use Drupal\Core\Session\AccountInterface;
-use Drupal\user\EntityOwnerInterface;
+use Drupal\Core\Session\AccountProxyInterface;
+use Symfony\Component\DependencyInjection\ContainerInterface;
 
 /**
  * Provides a block for creating a new content entity.
@@ -13,70 +23,161 @@ use Drupal\user\EntityOwnerInterface;
  * @Block(
  *   id = "entityform_block",
  *   admin_label = @Translation("Entity form"),
- *   category = @Translation("Forms")
+ *   category = @Translation("Forms"),
+ *   deriver = "Drupal\entityform_block\Plugin\Deriver\EntityFormDeriver"
  * )
  */
-class EntityEditFormBlock extends BlockBase {
+class EntityEditFormBlock extends BlockBase implements ContextAwarePluginInterface, ContainerFactoryPluginInterface {
+
+  /**
+   * The entity display repository.
+   *
+   * @var \Drupal\Core\Entity\EntityDisplayRepositoryInterface
+   */
+  protected $entityDisplayRepository;
+
+  /**
+   * The entity type bundle info service.
+   *
+   * @var \Drupal\Core\Entity\EntityTypeBundleInfoInterface
+   */
+  protected $entityTypeBundleInfo;
+
+  /**
+   * The entity type definition.
+   *
+   * @var \Drupal\Core\Entity\EntityTypeInterface|null
+   */
+  protected $entityType;
+
+  /**
+   * The entity type manager.
+   *
+   * @var \Drupal\Core\Entity\EntityTypeManagerInterface
+   */
+  protected $entityTypeManager;
+
+  /**
+   * The form builder.
+   *
+   * @var \Drupal\Core\Form\FormBuilder
+   */
+  protected $formBuilder;
+
+  /**
+   * The current user.
+   *
+   * @var \Drupal\Core\Session\AccountProxyInterface
+   */
+  protected $currentUser;
+
+  /**
+   * Creates a new EntityBundle instance.
+   *
+   * @param array $configuration
+   *   A configuration array containing information about the plugin instance.
+   * @param string $plugin_id
+   *   The plugin ID for the plugin instance.
+   * @param mixed $plugin_definition
+   *   The plugin implementation definition.
+   * @param \Drupal\Core\Entity\EntityDisplayRepositoryInterface $entity_display_repository
+   *   The entity display repository.
+   * @param \Drupal\Core\Entity\EntityTypeManagerInterface $entity_type_manager
+   *   The entity type manager.
+   * @param \Drupal\Core\Entity\EntityTypeBundleInfoInterface $entity_type_bundle_info
+   *   The entity type bundle info service.
+   * @param \Drupal\Core\Entity\EntityFormBuilder
+   *   The entity form builder.
+   */
+  public function __construct(array $configuration, $plugin_id, $plugin_definition, EntityDisplayRepositoryInterface $entity_display_repository, EntityTypeManagerInterface $entity_type_manager, EntityTypeBundleInfoInterface $entity_type_bundle_info, FormBuilder $form_builder, AccountProxyInterface $current_user) {
+    parent::__construct($configuration, $plugin_id, $plugin_definition);
+    $this->entityDisplayRepository = $entity_display_repository;
+    $this->entityTypeBundleInfo = $entity_type_bundle_info;
+    $this->entityTypeManager = $entity_type_manager;
+    $definition = $entity_type_manager->getDefinition($this->getDerivativeId());
+    if ($definition->entityClassImplements(FieldableEntityInterface::class)) {
+      $this->entityType = $definition;
+    }
+    else {
+      $this->entityType = $entity_type_manager->getDefinition($definition->getBundleOf());
+    }
+    $this->formBuilder = $form_builder;
+    $this->currentUser = $current_user;
+  }
 
   /**
    * {@inheritdoc}
    */
-  public function defaultConfiguration() {
-    return array(
-      'entity_type' => '',
-      'bundle' => '',
+  public static function create(ContainerInterface $container, array $configuration, $plugin_id, $plugin_definition) {
+    return new static(
+      $configuration,
+      $plugin_id,
+      $plugin_definition,
+      $container->get('entity_display.repository'),
+      $container->get('entity_type.manager'),
+      $container->get('entity_type.bundle.info'),
+      $container->get('form_builder'),
+      $container->get('current_user')
     );
   }
 
   /**
    * {@inheritdoc}
    */
+  public function defaultConfiguration() {
+    return [
+      'bundle' => NULL,
+      'view_mode' => 'default',
+    ];
+  }
+
+  /**
+   * {@inheritdoc}
+   */
   protected function blockAccess(AccountInterface $account) {
-    return \Drupal::entityManager()
-      ->getAccessControlHandler($this->configuration['entity_type'])
-      ->createAccess($this->configuration['bundle'], $account, [], TRUE);
+    $access_handler = $this->entityTypeManager
+      ->getAccessControlHandler($this->entityType->id());
+    /* @var $entity \Drupal\Core\Entity\EntityInterface */
+    $entity = $this->getContextValue('entity');
+    // If it is not a fieldable entity then it is an add form page so create a
+    // new entity.
+    if (!$entity instanceof FieldableEntityInterface || $entity->isNew()) {
+      return $access_handler
+        ->createAccess($this->configuration['bundle'], $account, [], TRUE);
+    }
+    return $access_handler->access($entity, 'update', $account, TRUE);
   }
 
   /**
    * {@inheritdoc}
    */
   public function blockForm($form, FormStateInterface $form_state) {
-    $form = parent::blockForm($form, $form_state);
-
-    // Get content entity types.
-    /** @var \Drupal\Core\Entity\ContentEntityTypeInterface[] $content_entity_types */
-    $content_entity_types = [];
-    foreach (\Drupal::entityManager()->getDefinitions() as $entity_type_id => $entity_type_definition) {
-      if ($entity_type_definition->getGroup() == 'content') {
-        $content_entity_types[$entity_type_id] = $entity_type_definition;
-      }
-    }
-    $options = array();
-    foreach ($content_entity_types as $type_key => $entity_type) {
-      // Entities that do not declare a form class.
-      // Exclude Comment entities as they have to be attached to another entity.
-      if (!$entity_type->hasFormClasses() || $type_key == 'comment') {
-        continue;
-      }
-      // Get all bundles for current entity type.
-      $entity_type_bundles = \Drupal::entityManager()->getBundleInfo($type_key);
-      foreach ($entity_type_bundles as $bundle_key => $bundle_info) {
+    $options = [];
+    $entity_type_id = $this->entityType->id();
+    if ($this->entityType->hasKey('bundle')) {
+      $bundles = $this->entityTypeBundleInfo->getBundleInfo($entity_type_id);
+      foreach ($bundles as $id => $info) {
         // Personal contact form requires a user recipient to be specified.
-        if ($bundle_key == 'personal' && $type_key == 'contact_message') {
+        if ($id == 'personal' && $entity_type_id == 'contact_message') {
           continue;
         }
-        $options[(string) $entity_type->getLabel()][$type_key . '.' . $bundle_key] = $bundle_info['label'];
+        $options[$id] = $info['label'];
       }
+
+      $form['bundle'] = [
+        '#title' => $this->t('Select the bundle'),
+        '#type' => 'select',
+        '#options' => $options,
+        '#default_value' => $this->configuration['bundle'],
+      ];
     }
 
-    $form['entity_type_bundle'] = array(
-      '#title' => $this->t('Entity Type + Bundle'),
+    $form['view_mode'] = [
       '#type' => 'select',
-      '#options' => $options,
-      '#required' => TRUE,
-      '#default_value' => $this->configuration['entity_type'] . '.' . $this->configuration['bundle'],
-    );
-
+      '#options' => $this->entityDisplayRepository->getFormModeOptions($entity_type_id),
+      '#title' => $this->t('View mode'),
+      '#default_value' => $this->configuration['view_mode'],
+    ];
     return $form;
   }
 
@@ -84,34 +185,31 @@ class EntityEditFormBlock extends BlockBase {
    * {@inheritdoc}
    */
   public function blockSubmit($form, FormStateInterface $form_state) {
-    $selected_entity_type_bundle = $form_state->getValue('entity_type_bundle');
-    $values = explode('.', $selected_entity_type_bundle);
-    $this->configuration['entity_type'] = $values[0];
-    $this->configuration['bundle'] = $values[1];
+    $this->configuration['bundle'] = $form_state->getValue('bundle');
+    $this->configuration['view_mode'] = $form_state->getValue('view_mode');
   }
 
   /**
    * {@inheritdoc}
    */
   public function build() {
-    $values = array();
-    // Specify selected bundle if the entity has bundles.
-    if (\Drupal::entityManager()->getDefinition($this->configuration['entity_type'])->hasKey('bundle')) {
-      $bundle_key = \Drupal::entityManager()->getDefinition($this->configuration['entity_type'])->getKey('bundle');
-      $values = array($bundle_key => $this->configuration['bundle']);
-    }
-
-    $entity = \Drupal::entityManager()
-      ->getStorage($this->configuration['entity_type'])
-      ->create($values);
+    /* @var $entity \Drupal\Core\Entity\EntityInterface */
+    $entity = $this->getContextValue('entity');
 
-    if ($entity instanceof EntityOwnerInterface) {
-      $entity->setOwnerId(\Drupal::currentUser()->id());
+    // If it is not a fieldable entity then it is an add form page so create a
+    // new entity.
+    if (!$entity instanceof FieldableEntityInterface || $entity->isNew()) {
+      $entity = $this->entityTypeManager->getStorage($this->entityType->id())
+        ->create([$this->entityType->getKey('bundle') => $this->configuration['bundle']]);
     }
-
-    $form = \Drupal::entityManager()
-      ->getFormObject($this->configuration['entity_type'], 'default')
+    $form_display = EntityFormDisplay::collectRenderDisplay($entity, $this->configuration['view_mode']);
+    // Set the form_display in form state.
+    $form_state = (new FormState())->set('form_display', $form_display);
+    // Get the default entity form object.
+    $form_object = $this->entityTypeManager->getFormObject($entity->getEntityTypeId(), 'default')
       ->setEntity($entity);
-    return \Drupal::formBuilder()->getForm($form);
+
+    return $this->formBuilder->buildForm($form_object, $form_state);
   }
+
 }
diff --git a/src/Plugin/Deriver/EntityFormDeriver.php b/src/Plugin/Deriver/EntityFormDeriver.php
new file mode 100644
index 0000000..5c84184
--- /dev/null
+++ b/src/Plugin/Deriver/EntityFormDeriver.php
@@ -0,0 +1,78 @@
+<?php
+
+namespace Drupal\entityform_block\Plugin\Deriver;
+
+use Drupal\Component\Plugin\Derivative\DeriverBase;
+use Drupal\Core\Entity\EntityTypeManagerInterface;
+use Drupal\Core\Plugin\Context\ContextDefinition;
+use Drupal\Core\Plugin\Discovery\ContainerDeriverInterface;
+use Drupal\Core\StringTranslation\StringTranslationTrait;
+use Drupal\Core\StringTranslation\TranslationInterface;
+use Symfony\Component\DependencyInjection\ContainerInterface;
+
+/**
+ * Provides entity form block definitions for each entity type.
+ */
+class EntityFormDeriver extends DeriverBase implements ContainerDeriverInterface {
+
+  use StringTranslationTrait;
+
+  /**
+   * The entity type manager.
+   *
+   * @var \Drupal\Core\Entity\EntityTypeManagerInterface
+   */
+  protected $entityTypeManager;
+
+  /**
+   * Constructs new EntityFormDeriver.
+   *
+   * @param \Drupal\Core\Entity\EntityTypeManagerInterface $entity_type_manager
+   *   The entity type manager.
+   * @param \Drupal\Core\StringTranslation\TranslationInterface $string_translation
+   *   The string translation service.
+   */
+  public function __construct(EntityTypeManagerInterface $entity_type_manager, TranslationInterface $string_translation) {
+    $this->entityTypeManager = $entity_type_manager;
+    $this->stringTranslation = $string_translation;
+  }
+
+  /**
+   * {@inheritdoc}
+   */
+  public static function create(ContainerInterface $container, $base_plugin_id) {
+    return new static(
+      $container->get('entity_type.manager'),
+      $container->get('string_translation')
+    );
+  }
+
+  /**
+   * {@inheritdoc}
+   */
+  public function getDerivativeDefinitions($base_plugin_definition) {
+    $entity_types = $this->entityTypeManager->getDefinitions();
+    foreach ($entity_types as $entity_type_id => $entity_type) {
+      // The content entity types that have a form class.
+      // Exclude Comment entities as they have to be attached to another entity.
+      if (($entity_type->getGroup() == 'content') && $entity_type->hasFormClasses() && $entity_type_id != 'comment') {
+        $this->derivatives[$entity_type_id] = $base_plugin_definition;
+        $this->derivatives[$entity_type_id]['admin_label'] = $this->t('Entity form (@label)', ['@label' => $entity_type->getLabel()]);
+        $this->derivatives[$entity_type_id]['context'] = [
+          'entity' => new ContextDefinition('entity:' . $entity_type_id),
+        ];
+      }
+      // If this entity is bundle of a content entity than add the context of
+      // bundle entity.
+      if (($bundle_of = $entity_type->getBundleOf()) && ($bundle_of_definition = $entity_types[$bundle_of]) && $bundle_of_definition->hasFormClasses() && $bundle_of != 'comment') {
+        $this->derivatives[$entity_type_id] = $base_plugin_definition;
+        $this->derivatives[$entity_type_id]['admin_label'] = $this->t('Entity form (@label)', ['@label' => $bundle_of_definition->getLabel()]);
+        $this->derivatives[$entity_type_id]['context'] = [
+          'entity' => new ContextDefinition('entity:' . $entity_type_id),
+        ];
+      }
+    }
+    return $this->derivatives;
+  }
+
+}