diff --git a/src/Plugin/Block/EntityEditFormBlock.php b/src/Plugin/Block/EntityEditFormBlock.php
index ab07fa5..eeeaaaf 100644
--- a/src/Plugin/Block/EntityEditFormBlock.php
+++ b/src/Plugin/Block/EntityEditFormBlock.php
@@ -8,6 +8,7 @@ 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;
@@ -94,7 +95,13 @@ class EntityEditFormBlock extends BlockBase implements ContextAwarePluginInterfa
     $this->entityDisplayRepository = $entity_display_repository;
     $this->entityTypeBundleInfo = $entity_type_bundle_info;
     $this->entityTypeManager = $entity_type_manager;
-    $this->entityType = $entity_type_manager->getDefinition($this->getDerivativeId());
+    $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;
   }
@@ -182,8 +189,9 @@ class EntityEditFormBlock extends BlockBase implements ContextAwarePluginInterfa
     /* @var $entity \Drupal\Core\Entity\EntityInterface */
     $entity = $this->getContextValue('entity');
 
-    // If it is not a content entity then add form page so create a new entity.
-    if (!$entity instanceof ContentEntityInterface) {
+    // If it is not a fieldable entity then it is an add form page so create a
+    // new entity.
+    if (!$entity instanceof FieldableEntityInterface) {
       $entity = $this->entityTypeManager->getStorage($this->entityType->id())
         ->create([$this->entityType->getKey('bundle') => $this->configuration['bundle']]);
     }
diff --git a/src/Plugin/Deriver/EntityFormDeriver.php b/src/Plugin/Deriver/EntityFormDeriver.php
index 140e67a..5c84184 100644
--- a/src/Plugin/Deriver/EntityFormDeriver.php
+++ b/src/Plugin/Deriver/EntityFormDeriver.php
@@ -65,9 +65,9 @@ class EntityFormDeriver extends DeriverBase implements ContainerDeriverInterface
       // 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[$bundle_of] = $base_plugin_definition;
-        $this->derivatives[$bundle_of]['admin_label'] = $this->t('Entity form (@label)', ['@label' => $bundle_of_definition->getLabel()]);
-        $this->derivatives[$bundle_of]['context'] = [
+        $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),
         ];
       }