diff --git a/wagtailmedia/wagtail_hooks.py b/wagtailmedia/wagtail_hooks.py
index 7387cb91..00777874 100644
--- a/wagtailmedia/wagtail_hooks.py
+++ b/wagtailmedia/wagtail_hooks.py
@@ -4,6 +4,7 @@
 from django.utils.translation import ugettext_lazy as _
 from django.utils.translation import ungettext
 
+from wagtail import VERSION as WAGTAIL_VERSION
 from wagtail.admin.menu import MenuItem
 from wagtail.admin.rich_text.editors.draftail import features as draftail_features
 from wagtail.admin.search import SearchArea
@@ -110,14 +111,16 @@ def describe_collection_media(collection):
 
 @hooks.register('register_rich_text_features')
 def register_custom_media_feature(features):
-    # Register a handler for converting <embed type="customimage"> to frontend HTML
-    features.register_embed_type(MediaEmbedHandler)
+	# Register a handler for converting <embed type="customimage"> to frontend HTML
+	if WAGTAIL_VERSION < (2, 5):
+		features.register_embed_type(embed_type='media', handler=MediaEmbedHandler)
+	else:
+		features.register_embed_type(handler=MediaEmbedHandler)
 
     feature_name = 'wagtailmedia'
-    type_ = 'MEDIA'
 
     control = {
-        'type': type_,
+        'type': 'MEDIA',
         'icon': 'media',
         'description': 'Media',
     }