diff --git a/entity_embed.module b/entity_embed.module
index ba1d1fd..4c76bc7 100644
--- a/entity_embed.module
+++ b/entity_embed.module
@@ -18,3 +18,54 @@ function entity_embed_help($route_name, RouteMatchInterface $route_match) {
       return '<p>' . t("Embed buttons define the buttons that can be added to CKEditor's toolbar. All buttons correspond to the selcted entity type, although a particular entity type can have more than one buttons. After you've created the desired buttons, go to <a href='@filterformats'>Text formats and editors configuration</a> page to add them to CKEditor's toolbar. Make sure to add <em>&lt;drupal-entity&gt;</em> to the list of 'Allowed HTML tags'.", array('@filterformats' => Url::fromRoute('filter.admin_overview')->toString())) . '</p>';
   }
 }
+
+/**
+ * Implements hook_embed_button_insert().
+ */
+function entity_embed_embed_button_insert(\Drupal\entity_embed\EmbedButtonInterface $entity) {
+  $button_icon_uuid = $entity->get('button_icon_uuid');
+  if ($button_icon_uuid) {
+    if ($file = \Drupal::entityManager()->loadEntityByUuid('file', $button_icon_uuid)) {
+      \Drupal::service('file.usage')->add($file, 'entity_embed', $entity->getEntityTypeId(), $entity->id());
+    }
+  }
+}
+
+/**
+ * Implements hook_embed_button_update().
+ */
+function entity_embed_embed_button_update(\Drupal\entity_embed\EmbedButtonInterface $entity) {
+  $new_button_icon_uuid = $entity->get('button_icon_uuid');
+
+  if (isset($entity->original)) {
+    $unchanged = $entity->original;
+    $old_button_icon_uuid = $unchanged->get('button_icon_uuid');
+
+    if (!empty($old_button_icon_uuid) && $old_button_icon_uuid != $new_button_icon_uuid) {
+      if ($file = \Drupal::entityManager()->loadEntityByUuid('file', $old_button_icon_uuid)) {
+        \Drupal::service('file.usage')->delete($file, 'entity_embed', $entity->getEntityTypeId(), $entity->id());
+      }
+    }
+  }
+
+  if ($new_button_icon_uuid) {
+    if ($file = \Drupal::entityManager()->loadEntityByUuid('file', $new_button_icon_uuid)) {
+      $usage = \Drupal::service('file.usage')->listUsage($file);
+      if (empty($usage['entity_embed'][$entity->getEntityTypeId()][$entity->id()])) {
+        \Drupal::service('file.usage')->add($file, 'entity_embed', $entity->getEntityTypeId(), $entity->id());
+      }
+    }
+  }
+}
+
+/**
+ * Implements hook_embed_button_update().
+ */
+function entity_embed_embed_button_delete(\Drupal\entity_embed\EmbedButtonInterface $entity) {
+  $button_icon_uuid = $entity->get('button_icon_uuid');
+  if ($button_icon_uuid) {
+    if ($file = \Drupal::entityManager()->loadEntityByUuid('file', $button_icon_uuid)) {
+      \Drupal::service('file.usage')->delete($file, 'entity_embed', $entity->getEntityTypeId(), $entity->id());
+    }
+  }
+}