diff --git a/src/Plugin/Field/FieldWidget/InlineEntityFormSimple.php b/src/Plugin/Field/FieldWidget/InlineEntityFormSimple.php
index ea7f266..ce47fde 100644
--- a/src/Plugin/Field/FieldWidget/InlineEntityFormSimple.php
+++ b/src/Plugin/Field/FieldWidget/InlineEntityFormSimple.php
@@ -34,7 +34,12 @@ class InlineEntityFormSimple extends InlineEntityFormBase {
     // WidgetSubmit will be needed once extractFormValues fills the $form_state.
     $parents = array_merge($element['#field_parents'], [$items->getName()]);
     $ief_id = sha1(implode('-', $parents));
-    $form_state->set(['inline_entity_form', $ief_id], []);
+
+    // Get state from storage.
+    $widget_state = $form_state->get(['inline_entity_form', $ief_id]);
+    if (!$widget_state){
+      $form_state->set(['inline_entity_form', $ief_id], []);
+    }
 
     $element = [
       '#type' => $this->getSetting('collapsible') ? 'details' : 'fieldset',
@@ -52,7 +57,14 @@ class InlineEntityFormSimple extends InlineEntityFormBase {
       $element['warning']['#markup'] = $this->t('Unable to load the referenced entity.');
       return $element;
     }
+
+    /** @var \Drupal\Core\Entity\EntityInterface $entity */
     $entity = $item->entity;
+
+    if (isset($widget_state['entities'][$delta]['entity'])) {
+      $entity = $widget_state['entities'][$delta]['entity'];
+    }
+
     $op = $entity ? 'edit' : 'add';
     $langcode = $items->getEntity()->language()->getId();
     $parents = array_merge($element['#field_parents'], [