diff --git a/wagtailmodelchoosers/client/DefaultDecorator.tsx b/wagtailmodelchoosers/client/DefaultDecorator.tsx
index 3333203..2ef5892 100644
--- a/wagtailmodelchoosers/client/DefaultDecorator.tsx
+++ b/wagtailmodelchoosers/client/DefaultDecorator.tsx
@@ -1,9 +1,10 @@
+import { ContentState } from "draft-js";
 import React from "react";
 
 type Props = {
-    entityKey: string;
-    contentState: any;
-    children: React.ReactNode;
+    entityKey?: string;
+    contentState: ContentState;
+    children?: React.ReactNode;
 };
 
 const DefaultDecorator = (entityData: any, props: Props) => {
diff --git a/wagtailmodelchoosers/client/draftailmodelchoosers.js b/wagtailmodelchoosers/client/draftailmodelchoosers.js
deleted file mode 100644
index 7b403ad..0000000
--- a/wagtailmodelchoosers/client/draftailmodelchoosers.js
+++ /dev/null
@@ -1,31 +0,0 @@
-import ModelSource from './sources/ModelSource';
-import DefaultDecorator from './DefaultDecorator';
-
-const WrapDecorator = (entityData, decorator) => (...args) => decorator(entityData, ...args);
-
-const modelChooserDraftailInit = (draftailOptions, modelChooserEntityTypes, constructor) => {
-  // Save entities for decorators to use.  Doing this because I can't find a way to get entity type
-  // data from within a decorator without saving it to contentstate, which ends up getting saved
-  // to DB, which is unnecessary and breaks historical data.
-  const modelChooserDraftailEntities = Object.fromEntries(
-    draftailOptions.entityTypes
-      .filter(et => modelChooserEntityTypes.includes(et.type))
-      .map(et => [et.type, et]),
-  );
-
-  // Register all modelchooser entity plugins.
-  modelChooserEntityTypes.forEach((entityType) => {
-    const decorators = window.draftailDecorators || {};
-    const decorator = decorators[entityType] || DefaultDecorator;
-    const plugin = {
-      type: entityType,
-      source: ModelSource,
-      decorator: WrapDecorator(modelChooserDraftailEntities[entityType], decorator),
-    };
-    window.draftail.registerPlugin(plugin);
-  });
-
-  return new window.telepath.constructors[constructor](draftailOptions);
-};
-
-window.telepath.register('modelChooserDraftailInit', modelChooserDraftailInit);
diff --git a/wagtailmodelchoosers/client/draftailmodelchoosers.tsx b/wagtailmodelchoosers/client/draftailmodelchoosers.tsx
new file mode 100644
index 0000000..544827e
--- /dev/null
+++ b/wagtailmodelchoosers/client/draftailmodelchoosers.tsx
@@ -0,0 +1,54 @@
+import ModelSource from "./sources/ModelSource";
+import DefaultDecorator from "./DefaultDecorator";
+
+const WrapDecorator =
+    (entityData, decorator) =>
+    (...args) =>
+        decorator(entityData, ...args);
+
+type Props = {
+    modelChooserEntityTypes: any[];
+    draftailOptions: any;
+    widgetAttrIds: any;
+};
+
+const modelChooserDraftailInit = ({
+    modelChooserEntityTypes,
+    draftailOptions,
+    widgetAttrIds,
+}: Props) => {
+    // Save entities for decorators to use.  Doing this because I can't find a way to get entity type
+    // data from within a decorator without saving it to contentstate, which ends up getting saved
+    // to DB, which is unnecessary and breaks historical data.
+
+    // Typescript doesn't know about .fromEntries, but it has been polyfilled for IE.
+    // @ts-ignore
+    const modelChooserDraftailEntities = Object.fromEntries(
+        draftailOptions.entityTypes
+            .filter((et) => modelChooserEntityTypes.includes(et.type))
+            .map((et) => [et.type, et])
+    );
+
+    // Register all modelchooser entity plugins.
+    modelChooserEntityTypes.forEach((entityType) => {
+        const decorators = (window as any).draftailDecorators || {};
+        const decorator = decorators[entityType] || DefaultDecorator;
+        const plugin = {
+            type: entityType,
+            source: ModelSource,
+            decorator: WrapDecorator(
+                modelChooserDraftailEntities[entityType],
+                decorator
+            ),
+        };
+        (window as any).draftail.registerPlugin(plugin);
+    });
+
+    (window as any).draftail.initEditor(
+        widgetAttrIds,
+        draftailOptions,
+        document.currentScript
+    );
+};
+
+(window as any).modelChooserDraftailInit = modelChooserDraftailInit;
diff --git a/wagtailmodelchoosers/client/sources/ModelSource.tsx b/wagtailmodelchoosers/client/sources/ModelSource.tsx
index 6c994e5..536b20a 100644
--- a/wagtailmodelchoosers/client/sources/ModelSource.tsx
+++ b/wagtailmodelchoosers/client/sources/ModelSource.tsx
@@ -1,14 +1,14 @@
 import React from "react";
 import "../polyfills";
-import { Modifier, EditorState } from "draft-js";
+import { Modifier, EditorState, Entity } from "draft-js";
 import ModelPicker from "../components/ModelPicker";
 
 const API_BASE_URL = "/admin/modelchoosers/api/v1";
 
 type Props = {
-    editorState: any;
+    editorState: EditorState;
     options: any;
-    entity?: any;
+    entity?: Entity;
     onClose: () => void;
     onUpdate: () => void;
     onComplete: (nextState: any) => void;
diff --git a/wagtailmodelchoosers/client/webpack.config.dev.js b/wagtailmodelchoosers/client/webpack.config.dev.js
index 1467c44..4a4ffab 100644
--- a/wagtailmodelchoosers/client/webpack.config.dev.js
+++ b/wagtailmodelchoosers/client/webpack.config.dev.js
@@ -9,7 +9,7 @@ module.exports = {
     wagtailmodelchoosers:
       "./wagtailmodelchoosers/client/wagtailmodelchoosers.tsx",
     draftailmodelchoosers:
-      "./wagtailmodelchoosers/client/draftailmodelchoosers.js",
+      "./wagtailmodelchoosers/client/draftailmodelchoosers.tsx",
     polyfills: "./wagtailmodelchoosers/client/polyfills.js",
   },
   output: {