diff --git a/includes/entity_embed.display.inc b/includes/entity_embed.display.inc
index f250582..984a3a2 100644
--- a/includes/entity_embed.display.inc
+++ b/includes/entity_embed.display.inc
@@ -55,17 +55,15 @@ function entity_embed_render_entity($entity_type, $entity, array $context = arra
     $context
   );
 
-  // Maintain data-align if it is there.
+  // Maintain data-align.
   if (isset($context['data-align'])) {
-    $build['#attributes']['data-align'] = $context['data-align'];
-  }
-  elseif ((isset($context['class']))) {
-    $build['#attributes']['class'][] = $context['class'];
-  }
-
-  // Maintain data-caption if it is there.
-  if (isset($context['data-caption'])) {
-    $build['#attributes']['data-caption'] = $context['data-caption'];
+    if (isset($build[$entity_type])) {
+      foreach ($build[$entity_type] as $entity_id => $entity_array) {
+        if (is_array($entity_array)) {
+          $build[$entity_type][$entity_id]['#align'] = $context['data-align'];
+        }
+      }
+    }
   }
 
   // @todo Should this hook get invoked if $build is an empty array?
@@ -76,6 +74,17 @@ function entity_embed_render_entity($entity_type, $entity, array $context = arra
 }
 
 /**
+ * Implements hook_preprocess_HOOK.
+ *
+ * Add the alignment class to the existing div wrapper.
+ */
+function entity_embed_preprocess_file_entity(&$variables) {
+  if(isset($variables['elements']['#align'])) {
+    $variables['classes_array'][] = 'align-' .$variables['elements']['#align'];
+  }
+}
+
+/**
  * Renders an entity using an Entity Embed Display plugin.
  *
  * @param $entity_type