diff --git a/remix/app/components/ImageVaultDAM.tsx b/remix/app/components/ImageVaultDAM.tsx
index 196264c..8b49d25 100644
--- a/remix/app/components/ImageVaultDAM.tsx
+++ b/remix/app/components/ImageVaultDAM.tsx
@@ -7,16 +7,12 @@ import {
   FieldLabel,
   cbModal,
 } from "@contentstack/venus-components"
-import {
-  getImageVaultAssetFromData,
-  isInsertResponse,
-  openImageVault,
-} from "~/utils/imagevault"
+import { isInsertResponse, openImageVault } from "~/utils/imagevault"
 import ImageEditModal from "./ImageEditModal"
 
 import type UiLocation from "@contentstack/app-sdk/dist/src/uiLocation"
 import type { CbModalProps } from "@contentstack/venus-components/build/components/Modal/Modal"
-import type { ImageVaultAsset, InsertResponse } from "~/types/imagevault"
+import type { ImageVaultAsset } from "~/types/imagevault"
 import type { Lang } from "~/types/lang"
 import type {
   EntryDataPublishDetails,
@@ -27,7 +23,7 @@ import FullSizeImage from "./FullSizeImage"
 export type ImageVaultDAMProps = {
   sdk: UiLocation
   config: ImageVaultDAMConfig
-  initialData: ImageVaultAsset | InsertResponse | null
+  imageVaultAsset: ImageVaultAsset | null
 }
 
 type DAMButtonProps = { onClick: () => void }
@@ -115,9 +111,8 @@ function Media({ media, onDelete, onEdit }: MediaProps) {
 export default function ImageVaultDAM({
   sdk,
   config,
-  initialData,
+  imageVaultAsset,
 }: ImageVaultDAMProps) {
-  const imageVaultAsset = getImageVaultAssetFromData(initialData)
   const [media, setMedia] = useState(imageVaultAsset)
 
   const field = sdk.location.CustomField?.field
@@ -176,18 +171,16 @@ export default function ImageVaultDAM({
     updateFrameHeight()
   }, [updateFrameHeight])
 
-  // The existing data might still be in InsertResponse format if the user has not edited it yet.
-  // We'll convert it to ImageVaultAsset when the user edits the the emtry.
-  useEffect(() => {
-    if (isInsertResponse(initialData) && imageVaultAsset) {
-      handleMedia(imageVaultAsset)
-    }
-  }, [initialData, imageVaultAsset, handleMedia])
-
   if (!field || !frame || !entry || !stack) {
     return <p>Initializing custom field...</p>
   }
 
+  const fieldData = field.getData()
+
+  if (isInsertResponse(fieldData)) {
+    field.setData(imageVaultAsset)
+  }
+
   const entryData: EntryDataPublishDetails = {
     //TODO: Add support for branches
     branch: "main",
diff --git a/remix/app/routes/field.tsx b/remix/app/routes/field.tsx
index 36ed3db..1419669 100644
--- a/remix/app/routes/field.tsx
+++ b/remix/app/routes/field.tsx
@@ -9,8 +9,11 @@ import InvalidConfig from "~/components/InvalidConfig"
 
 import { GenericObjectType } from "@contentstack/app-sdk/dist/src/types/common.types"
 import UiLocation from "@contentstack/app-sdk/dist/src/uiLocation"
-import type { InsertResponse } from "~/types/imagevault"
-import { isImageVaultDAMConfig } from "~/utils/imagevault"
+import type { ImageVaultAsset, InsertResponse } from "~/types/imagevault"
+import {
+  getImageVaultAssetFromData,
+  isImageVaultDAMConfig,
+} from "~/utils/imagevault"
 
 const ImageVaultDAM = lazy(() => import("~/components/ImageVaultDAM"))
 
@@ -25,7 +28,9 @@ function FieldContent({ sdk, appConfig }: FieldContentProps) {
   )
 
   const [showDisclaimer, setShowDisclaimer] = useState(false)
-  const [fieldData, setFieldData] = useState<InsertResponse | null>()
+  const [fieldData, setFieldData] = useState<
+    InsertResponse | ImageVaultAsset | null
+  >()
   const [dataIsLoaded, setDataIsLoaded] = useState(false)
 
   const entry = sdk?.location.CustomField?.entry
@@ -76,9 +81,17 @@ function FieldContent({ sdk, appConfig }: FieldContentProps) {
   const fieldConfig = sdk.location.CustomField?.fieldConfig
   const config = { ...appConfig, ...fieldConfig }
 
+  // The existing data might still be in InsertResponse format if the user has not edited it yet.
+  // We'll convert it to ImageVaultAsset when the user edits the the entry.
+  const imageVaultAsset = getImageVaultAssetFromData(initialData)
+
   return (
     <Suspense fallback={<p>Loading field...</p>}>
-      <ImageVaultDAM config={config} sdk={sdk} initialData={initialData} />
+      <ImageVaultDAM
+        config={config}
+        sdk={sdk}
+        imageVaultAsset={imageVaultAsset}
+      />
     </Suspense>
   )
 }
diff --git a/rte/components/ImageElement.tsx b/rte/components/ImageElement.tsx
index 9c2065c..fc5c985 100644
--- a/rte/components/ImageElement.tsx
+++ b/rte/components/ImageElement.tsx
@@ -1,5 +1,5 @@
 import { Icon, Tooltip, cbModal } from "@contentstack/venus-components"
-import React, { PropsWithChildren, useCallback, useEffect } from "react"
+import React, { PropsWithChildren, useCallback } from "react"
 
 import EmbedBtn from "./EmbedBtn"
 import ImageEditModal from "./ImageEditModal"
@@ -19,6 +19,7 @@ type ImageElementProps = PropsWithChildren & {
   rte: IRteParam
 }
 export function ImageElement({ children, element, rte }: ImageElementProps) {
+  const isAssetInsertResponse = isInsertResponse(element.attrs)
   const imageVaultAsset = getImageVaultAssetFromData(element.attrs)
   const isSelected = rte.selection.isSelected()
   const isFocused = rte.selection.isFocused()
@@ -62,14 +63,6 @@ export function ImageElement({ children, element, rte }: ImageElementProps) {
     })
   }, [element, handleMedia])
 
-  // The existing data might still be in InsertResponse format if the user has not edited it yet.
-  // We'll convert it to ImageVaultAsset when the user edits the RTE.
-  useEffect(() => {
-    if (isInsertResponse(element.attrs) && imageVaultAsset) {
-      handleMedia(imageVaultAsset)
-    }
-  }, [element.attrs, imageVaultAsset, handleMedia])
-
   const ToolTipButtons = () => {
     return (
       <div contentEditable={false} className="embed--btn-group">
@@ -91,6 +84,11 @@ export function ImageElement({ children, element, rte }: ImageElementProps) {
   if (!imageVaultAsset) {
     return <>{children}</>
   }
+  // The existing data might still be in InsertResponse format if the user has not edited it yet.
+  // We'll convert it to ImageVaultAsset when the user edits the RTE.
+  if (isAssetInsertResponse) {
+    handleMedia(imageVaultAsset)
+  }
 
   return (
     <Tooltip