Draft-JS - Translating <img> into an atomic block with convertFromHTML gets rid of Entity

So I have been trying to get the method convertFromHTML to translate Images into an atomic block so that it may be compatible with the draft-js-image-plugin since it expects blocks of the type atomic

Given a simple HTML structure with some text and an image out of the box the convertFromHTML produces this contentState:

{
  "blocks": [
    {
      "key": "82k8",
      "text": "‐‐‐‐‐‐‐ Original Message ‐‐‐‐‐‐‐",
      "type": "unstyled",
      "depth": 0,
      "inlineStyleRanges": [],
      "entityRanges": [],
      "data": {}
    },
    {
      "key": "9jbor",
      "text": "On December 29, 2020, 5:20 PM EST  txwbi.nrjrtn@gmail.com wrote:",
      "type": "unstyled",
      "depth": 0,
      "inlineStyleRanges": [],
      "entityRanges": [{ "offset": 34, "length": 23, "key": 0 }],
      "data": {}
    },
    {
      "key": "anq8o",
      "text": "📷A bunch of text here to test out the body",
      "type": "unstyled",
      "depth": 0,
      "inlineStyleRanges": [
        { "offset": 3, "length": 13, "style": "ITALIC" },
        { "offset": 3, "length": 13, "style": "UNDERLINE" },
        { "offset": 38, "length": 4, "style": "BOLD" }
      ],
      "entityRanges": [{ "offset": 0, "length": 1, "key": 1 }],
      "data": {}
    }
  ],
  "entityMap": {
    "0": {
      "type": "LINK",
      "mutability": "MUTABLE",
      "data": {
        "href": "mailto:txwbi.nrjrtn@gmail.com",
        "rel": "noreferrer nofollow noopener",
        "target": "_blank",
        "url": "mailto:txwbi.nrjrtn@gmail.com"
      }
    },
    "1": {
      "type": "IMAGE",
      "mutability": "IMMUTABLE",
      "data": {
        "alt": "cory_emoji.png",
        "height": "210",
        "src": "data:image/png;base64, ...BASE64ENCODEDIMAGE WOULD BE HERE I REMOVED BECAUSE OF CHAR LIMITS",
        "width": "173"
      }
    }
  }
}

Which we can see that the img tag takes on a unstyled block which is not what I'd like. So I created the following function that extends the default Block Render Map to account for the img tag:

const {
  EditorState,
  convertToRaw,
  DefaultDraftBlockRenderMap,
  ContentState,
  convertFromHTML,
  getSafeBodyFromHTML
} = require('draft-js');

const Immutable = require('immutable');

module.exports.editorStateFromHTML = htmlBody => {
  console.log('HTML ---> EDITOR ::: RAW BODY', htmlBody);
  const blockRenderMap = Immutable.Map({
    atomic: {
      element: 'figure',
      aliasedElements: ['img']
    }
  });

  const extendedBlockRenderMap = DefaultDraftBlockRenderMap.merge(
    blockRenderMap
  );

  const blocksFromHTML = convertFromHTML(
    htmlBody,
    getSafeBodyFromHTML,
    extendedBlockRenderMap
  );

  console.log(blocksFromHTML);
  const state = ContentState.createFromBlockArray(
    blocksFromHTML.contentBlocks,
    blocksFromHTML.entityMap
  );
  console.log(JSON.stringify(convertToRaw(state)));
  const newEditor = EditorState.createWithContent(state);
  return newEditor;
};

Which results in this content State:

{
  "blocks": [
    {
      "key": "fhgqd",
      "text": "‐‐‐‐‐‐‐ Original Message ‐‐‐‐‐‐‐",
      "type": "unstyled",
      "depth": 0,
      "inlineStyleRanges": [],
      "entityRanges": [],
      "data": {}
    },
    {
      "key": "2nsnk",
      "text": "On December 29, 2020, 5:20 PM EST  txwbi.nrjrtn@gmail.com wrote:",
      "type": "unstyled",
      "depth": 0,
      "inlineStyleRanges": [],
      "entityRanges": [{ "offset": 34, "length": 23, "key": 0 }],
      "data": {}
    },
    {
      "key": "7c7cu",
      "text": "A bunch of text here to test out the body",
      "type": "unstyled",
      "depth": 0,
      "inlineStyleRanges": [
        { "offset": 2, "length": 13, "style": "UNDERLINE" },
        { "offset": 2, "length": 13, "style": "ITALIC" },
        { "offset": 37, "length": 4, "style": "BOLD" }
      ],
      "entityRanges": [],
      "data": {}
    },
    {
      "key": "f84vb",
      "text": "",
      "type": "atomic",
      "depth": 0,
      "inlineStyleRanges": [],
      "entityRanges": [],
      "data": {}
    }
  ],
  "entityMap": {
    "0": {
      "type": "LINK",
      "mutability": "MUTABLE",
      "data": {
        "href": "mailto:txwbi.nrjrtn@gmail.com",
        "rel": "noreferrer nofollow noopener",
        "target": "_blank",
        "url": "mailto:txwbi.nrjrtn@gmail.com"
      }
    }
  }
}

As you can see the entityMap now only has one key and all the data for the image is no longer there. How can I get it to make the img tag an atomic block while still creating the IMAGE entity in the entityMap???

My work around was has follow it's kind of ugly but it works and gives me the right output:

const {
  EditorState,
  convertToRaw,
  convertFromRaw,
  DefaultDraftBlockRenderMap,
  ContentState,
  convertFromHTML,
  getSafeBodyFromHTML
} = require('draft-js');

const Immutable = require('immutable');
const clone = require('rfdc')();

module.exports.editorStateFromHTML = htmlBody => {
  console.log('HTML ---> EDITOR ::: RAW BODY', htmlBody);
  const blockRenderMap = Immutable.Map({
    image: {
      element: 'img'
    }
  });

  const extendedBlockRenderMap = DefaultDraftBlockRenderMap.merge(
    blockRenderMap
  );

  const blocksFromHTML = convertFromHTML(
    htmlBody,
    getSafeBodyFromHTML,
    extendedBlockRenderMap
  );

  const state = ContentState.createFromBlockArray(
    blocksFromHTML.contentBlocks,
    blocksFromHTML.entityMap
  );
  const { blocks, entityMap } = convertToRaw(state);
  const imgCount = blocks.filter(b => b.type === 'image').length;

  if (imgCount > 0) {
    const fixedEntityMap = clone(entityMap);
    let arrKeys = Object.keys(fixedEntityMap);
    arrKeys = arrKeys.map(k => parseInt(k, 10));
    const lastKey = Math.max(...arrKeys);
    let blockCounter = lastKey;

    const fixedContentBlocks = blocks.map(blck => {
      if (blck.type === 'image') {
        blockCounter += 1;
        return {
          ...blck,
          text: ' ',
          type: 'atomic',
          entityRanges: [{ offset: 0, length: 1, key: blockCounter }]
        };
      }
      return blck;
    });

    const blocksFromHTML2 = convertFromHTML(htmlBody);
    const state2 = ContentState.createFromBlockArray(
      blocksFromHTML2.contentBlocks,
      blocksFromHTML2.entityMap
    );
    const { entityMap: imgEntities } = convertToRaw(state2);

    let entityCounter = lastKey;
    // eslint-disable-next-line no-restricted-syntax
    for (const [key, value] of Object.entries(imgEntities)) {
      if (value.type === 'IMAGE') {
        entityCounter += 1;
        fixedEntityMap[entityCounter] = value;
      }
    }

    const editorDefaultValue = {
      blocks: fixedContentBlocks,
      entityMap: fixedEntityMap
    };
    return EditorState.createWithContent(convertFromRaw(editorDefaultValue));
  }

  return EditorState.createWithContent(state);
};

If Anyone has a better solution I am all ears.