diff --git a/Model/Product/Gallery/CreateHandlerOverride.php b/Model/Product/Gallery/CreateHandler.php
similarity index 99%
rename from Model/Product/Gallery/CreateHandlerOverride.php
rename to Model/Product/Gallery/CreateHandler.php
index cc01f79..19fd05b 100644
--- a/Model/Product/Gallery/CreateHandlerOverride.php
+++ b/Model/Product/Gallery/CreateHandler.php
@@ -6,13 +6,13 @@
 use Magento\Catalog\Model\Product;
 use Magento\Framework\App\Filesystem\DirectoryList;
 use Magento\Framework\App\ObjectManager;
-use Magento\Framework\EntityManager\Operation\ExtensionInterface;
 use Magento\MediaStorage\Model\File\Uploader as FileUploader;
 use Magento\Store\Model\Store;
 use Magento\Store\Model\StoreManagerInterface;
 use Scaleflex\Filerobot\Model\FilerobotConfig;
+use Magento\Catalog\Model\Product\Gallery\CreateHandler as BaseCreateHandler;
 
-class CreateHandlerOverride implements ExtensionInterface
+class CreateHandler extends BaseCreateHandler
 {
     /**
      * @var \Magento\Framework\EntityManager\EntityMetadata
diff --git a/Model/Product/Gallery/Staging/CreateHandler.php b/Model/Product/Gallery/Staging/CreateHandler.php
new file mode 100644
index 0000000..4d41ba0
--- /dev/null
+++ b/Model/Product/Gallery/Staging/CreateHandler.php
@@ -0,0 +1,126 @@
+<?php
+
+namespace Scaleflex\Filerobot\Model\Product\Gallery\Staging;
+
+use Magento\Catalog\Api\Data\ProductInterface;
+use Magento\Catalog\Api\ProductAttributeRepositoryInterface;
+use Magento\Catalog\Model\Product;
+use Magento\Catalog\Model\Product\Gallery\CopyHandler;
+use Scaleflex\Filerobot\Model\FilerobotConfig;
+use Scaleflex\Filerobot\Model\Product\Gallery\CreateHandler as ProductGalleryCreateHandler;
+use Scaleflex\Filerobot\Model\Product\Gallery\UpdateHandler;
+use Magento\Catalog\Model\Product\Media\Config;
+use Magento\Catalog\Model\ResourceModel\Product\Gallery;
+use Magento\Framework\App\ObjectManager;
+use Magento\Framework\EntityManager\MetadataPool;
+use Magento\Framework\EntityManager\Operation\ExtensionInterface;
+use Magento\Framework\Exception\LocalizedException;
+use Magento\Framework\Filesystem;
+use Magento\Framework\Json\Helper\Data;
+use Magento\MediaStorage\Helper\File\Storage\Database;
+use Magento\Staging\Model\ResourceModel\Db\ReadEntityVersion;
+use Magento\Staging\Model\VersionHistoryInterface;
+use Magento\Staging\Model\VersionManager;
+
+
+class CreateHandler extends ProductGalleryCreateHandler implements ExtensionInterface
+{
+    /**
+     * @var UpdateHandler
+     */
+    private $updateHandler;
+
+    /**
+     * @var CopyHandler
+     */
+    private $copyHandler;
+
+    /**
+     * @var VersionHistoryInterface
+     */
+    private $versionHistory;
+
+    /**
+     * @var ReadEntityVersion
+     */
+    private $readEntityVersion;
+
+    /**
+     * @param MetadataPool $metadataPool
+     * @param ProductAttributeRepositoryInterface $attributeRepository
+     * @param Gallery $resourceModel
+     * @param Data $jsonHelper
+     * @param Config $mediaConfig
+     * @param Filesystem $filesystem
+     * @param Database $fileStorageDb
+     * @param UpdateHandler $updateHandler
+     * @param CopyHandler|null $copyHandler
+     * @param VersionHistoryInterface|null $versionHistory
+     * @param ReadEntityVersion|null $readEntityVersion
+     * @SuppressWarnings(PHPMD.ExcessiveParameterList)
+     */
+    public function __construct(
+        MetadataPool $metadataPool,
+        ProductAttributeRepositoryInterface $attributeRepository,
+        Gallery $resourceModel,
+        Data $jsonHelper,
+        Config $mediaConfig,
+        Filesystem $filesystem,
+        Database $fileStorageDb,
+        UpdateHandler $updateHandler,
+        FilerobotConfig $filerobotConfig,
+        ?CopyHandler $copyHandler = null,
+        ?VersionHistoryInterface $versionHistory = null,
+        ?ReadEntityVersion $readEntityVersion = null
+    ) {
+        $this->updateHandler = $updateHandler;
+
+        parent::__construct(
+            $metadataPool,
+            $attributeRepository,
+            $resourceModel,
+            $jsonHelper,
+            $mediaConfig,
+            $filesystem,
+            $fileStorageDb,
+            $filerobotConfig,
+        );
+        $this->copyHandler = $copyHandler ?? ObjectManager::getInstance()->get(CopyHandler::class);
+        $this->versionHistory = $versionHistory ?? ObjectManager::getInstance()->get(VersionHistoryInterface::class);
+        $this->readEntityVersion = $readEntityVersion ?? ObjectManager::getInstance()->get(ReadEntityVersion::class);
+    }
+
+    /**
+     * Execute create handler
+     *
+     * @param Product $product
+     * @param array $arguments
+     * @return bool|object
+     * @throws LocalizedException
+     */
+    public function execute($product, $arguments = [])
+    {
+        if (isset($arguments['origin_in'])) {
+            $originId = $arguments['origin_in'];
+        } elseif (isset($arguments['copy_origin_in'])) {
+            $originId = $arguments['copy_origin_in'];
+        } else {
+            $originId = $this->versionHistory->getCurrentId();
+        }
+
+        if (!in_array($product->getData('created_in'), [VersionManager::MIN_VERSION, $originId], true)) {
+            $arguments['original_link_id'] = $this->readEntityVersion->getVersionRowId(
+                ProductInterface::class,
+                $product->getData($this->metadata->getIdentifierField()),
+                $originId
+            );
+            $this->copyHandler->execute($product, $arguments);
+            if (!empty($arguments['is_rollback'])) {
+                return $product;
+            }
+        }
+        return $product->isObjectNew()
+            ? parent::execute($product, $arguments)
+            : $this->updateHandler->execute($product, $arguments);
+    }
+}
diff --git a/Model/Product/Gallery/Staging/UpdateHandler.php b/Model/Product/Gallery/Staging/UpdateHandler.php
new file mode 100644
index 0000000..d2bdfb7
--- /dev/null
+++ b/Model/Product/Gallery/Staging/UpdateHandler.php
@@ -0,0 +1,147 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Scaleflex\Filerobot\Model\Product\Gallery\Staging;
+
+use Magento\Catalog\Api\Data\ProductInterface;
+use Magento\Catalog\Api\ProductAttributeRepositoryInterface;
+use Magento\Catalog\Model\Product;
+use Magento\Catalog\Model\Product\Gallery\CopyHandler;
+use Magento\Catalog\Model\Product\Gallery\DeleteHandler;
+use Magento\Catalog\Model\Product\Gallery\ReadHandler;
+use Magento\Catalog\Model\Product\Media\Config;
+use Magento\Catalog\Model\ResourceModel\Product\Gallery;
+use Magento\Eav\Model\ResourceModel\AttributeValue;
+use Magento\Framework\EntityManager\MetadataPool;
+use Magento\Framework\Filesystem;
+use Magento\Framework\Json\Helper\Data;
+use Magento\MediaStorage\Helper\File\Storage\Database;
+use Magento\Staging\Model\ResourceModel\Db\ReadEntityVersion;
+use Magento\Store\Model\StoreManagerInterface;
+use Scaleflex\Filerobot\Model\FilerobotConfig;
+use Scaleflex\Filerobot\Model\Product\Gallery\UpdateHandler as BaseUpdateHandler;
+
+class UpdateHandler extends BaseUpdateHandler
+{
+    /**
+     * @var CopyHandler
+     */
+    private $copyHandler;
+
+    /**
+     * @var DeleteHandler
+     */
+    private $deleteHandler;
+
+    /**
+     * @var ReadHandler
+     */
+    private $readHandler;
+
+    /**
+     * @var ReadEntityVersion
+     */
+    private $readEntityVersion;
+
+    /**
+     * @var string[]
+     */
+    private $mediaAttributesWithLabels = [
+        'image',
+        'small_image',
+        'thumbnail'
+    ];
+
+    /**
+     * @param MetadataPool $metadataPool
+     * @param ProductAttributeRepositoryInterface $attributeRepository
+     * @param Gallery $resourceModel
+     * @param Data $jsonHelper
+     * @param Config $mediaConfig
+     * @param Filesystem $filesystem
+     * @param Database $fileStorageDb
+     * @param StoreManagerInterface $storeManager
+     * @param AttributeValue $attributeValue
+     * @param CopyHandler $copyHandler
+     * @param DeleteHandler $deleteHandler
+     * @param ReadHandler $readHandler
+     * @param ReadEntityVersion $readEntityVersion
+     * @SuppressWarnings(PHPMD.ExcessiveParameterList)
+     */
+    public function __construct(
+        MetadataPool $metadataPool,
+        ProductAttributeRepositoryInterface $attributeRepository,
+        Gallery $resourceModel,
+        Data $jsonHelper,
+        Config $mediaConfig,
+        Filesystem $filesystem,
+        Database $fileStorageDb,
+        StoreManagerInterface $storeManager,
+        AttributeValue $attributeValue,
+        CopyHandler $copyHandler,
+        DeleteHandler $deleteHandler,
+        ReadHandler $readHandler,
+        ReadEntityVersion $readEntityVersion,
+        FilerobotConfig $filerobotConfig
+    ) {
+        parent::__construct(
+            $metadataPool,
+            $attributeRepository,
+            $resourceModel,
+            $jsonHelper,
+            $mediaConfig,
+            $filesystem,
+            $filerobotConfig,
+            $fileStorageDb,
+            $storeManager,
+            $attributeValue
+        );
+        $this->copyHandler = $copyHandler;
+        $this->deleteHandler = $deleteHandler;
+        $this->readHandler = $readHandler;
+        $this->readEntityVersion = $readEntityVersion;
+    }
+
+    /**
+     * Update product media gallery
+     *
+     * @param Product $product
+     * @param array $arguments
+     * @return object
+     */
+    public function execute($product, $arguments = [])
+    {
+        if (!empty($arguments['is_rollback']) && !empty($arguments['copy_origin_in'])) {
+            $arguments['media_attribute_codes'] = $this->getMediaAttributeCodes();
+            $arguments['original_link_id'] = $this->readEntityVersion->getVersionRowId(
+                ProductInterface::class,
+                $product->getData($this->metadata->getIdentifierField()),
+                $arguments['copy_origin_in']
+            );
+            $this->deleteHandler->execute($product, $arguments);
+            $this->copyHandler->execute($product, $arguments);
+            // reload gallery data as new value IDs are auto generated
+            $this->readHandler->execute($product, $arguments);
+            return $product;
+        }
+        return parent::execute($product, $arguments);
+    }
+
+    /**
+     * Get all media attributes codes including their corresponding labels
+     *
+     * @return array
+     */
+    private function getMediaAttributeCodes(): array
+    {
+        $attributeCodes = [];
+        foreach ($this->mediaConfig->getMediaAttributeCodes() as $attributeCode) {
+            $attributeCodes[] = $attributeCode;
+            if (in_array($attributeCode, $this->mediaAttributesWithLabels)) {
+                $attributeCodes[] = $attributeCode . '_label';
+            }
+        }
+        return $attributeCodes;
+    }
+}
diff --git a/Model/Product/Gallery/UpdateHandlerOverride.php b/Model/Product/Gallery/UpdateHandler.php
similarity index 98%
rename from Model/Product/Gallery/UpdateHandlerOverride.php
rename to Model/Product/Gallery/UpdateHandler.php
index a3a59ca..31d17f4 100644
--- a/Model/Product/Gallery/UpdateHandlerOverride.php
+++ b/Model/Product/Gallery/UpdateHandler.php
@@ -6,7 +6,6 @@
 use Magento\Catalog\Api\ProductAttributeRepositoryInterface;
 use Magento\Catalog\Model\Product;
 use Scaleflex\Filerobot\Model\FilerobotConfig;
-use Scaleflex\Filerobot\Model\Product\Gallery\CreateHandlerOverride;
 use Magento\Catalog\Model\Product\Media\Config;
 use Magento\Catalog\Model\ResourceModel\Product\Gallery;
 use Magento\Eav\Model\ResourceModel\AttributeValue;
@@ -18,7 +17,7 @@
 use Magento\Store\Model\Store;
 use Magento\Store\Model\StoreManagerInterface;
 
-class UpdateHandlerOverride extends CreateHandlerOverride
+class UpdateHandler extends CreateHandler
 {
     /**
      * @var AttributeValue
diff --git a/Plugin/AddImageToGallery.php b/Plugin/AddImageToGallery.php
index 2fd102f..d306ce0 100644
--- a/Plugin/AddImageToGallery.php
+++ b/Plugin/AddImageToGallery.php
@@ -42,7 +42,7 @@ public function afterGetGalleryImages(Gallery $gallery, $images)
         foreach ($images as $image) {
             if ($this->fileRobotConfig->isFilerobot($image->getData('file'))) {
                 $url = $image->getData('file');
-                $entityId = $image->getData('entity_id');
+                $entityId = $image->getData('entity_id') ? $image->getData('entity_id') : $image->getData('row_id');
                 $product  = $this->productRepository->getById($entityId);
                 if ($product) {
                     $thumbImageSize = $this->imageHelper->init($product, 'product_thumbnail_image');
diff --git a/Ui/DataProvider/Product/Form/Modifier/Related.php b/Ui/DataProvider/Product/Form/Modifier/Related.php
index 97b6d13..b3cff2f 100644
--- a/Ui/DataProvider/Product/Form/Modifier/Related.php
+++ b/Ui/DataProvider/Product/Form/Modifier/Related.php
@@ -38,19 +38,21 @@ public function afterModifyData(
     ) {
         foreach ($result as $key => $value) {
             if ($key !== 'config') {
-                $links = $value['links'];
-                foreach ($links as &$link) {
-                    if (!empty($link)) {
-                        foreach ($link as &$linked) {
-                            $product = $this->productRepository->getById($linked['id']);
-                            $images = $product->getMediaAttributeValues();
-                            if (!empty($images) && $images['thumbnail'] && $this->fileRobotConfig->isFilerobot($images['thumbnail'])) {
-                                $linked['thumbnail'] = $images['thumbnail'];
+                if (array_key_exists('links', $value)) {
+                    $links = $value['links'];
+                    foreach ($links as &$link) {
+                        if (!empty($link)) {
+                            foreach ($link as &$linked) {
+                                $product = $this->productRepository->getById($linked['id']);
+                                $images = $product->getMediaAttributeValues();
+                                if (!empty($images) && $images['thumbnail'] && $this->fileRobotConfig->isFilerobot($images['thumbnail'])) {
+                                    $linked['thumbnail'] = $images['thumbnail'];
+                                }
                             }
                         }
                     }
+                    $result[$key]['links'] = $links;
                 }
-                $result[$key]['links'] = $links;
             }
         }
 
diff --git a/composer.json b/composer.json
index 704ba9f..9960cef 100644
--- a/composer.json
+++ b/composer.json
@@ -1,6 +1,6 @@
 {
-    "name": "scaleflex/module-filerobot",
-    "description": "Scaleflex Filerobot Magento Extension",
+    "name": "scaleflex/module-filerobot-enterprise",
+    "description": "Scaleflex Filerobot Magento Enterprise Extension",
     "type": "magento2-module",
     "license": "GPL-3.0",
     "authors": [
diff --git a/etc/di.xml b/etc/di.xml
index c4a5442..0822765 100644
--- a/etc/di.xml
+++ b/etc/di.xml
@@ -1,9 +1,13 @@
 <?xml version="1.0"?>
 <config xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:noNamespaceSchemaLocation="urn:magento:framework:ObjectManager/etc/config.xsd">
        <preference for="Magento\Catalog\Model\Product\Gallery\CreateHandler"
-                    type="Scaleflex\Filerobot\Model\Product\Gallery\CreateHandlerOverride" />
+                    type="Scaleflex\Filerobot\Model\Product\Gallery\CreateHandler" />
        <preference for="Magento\Catalog\Model\Product\Gallery\UpdateHandler"
-                   type="Scaleflex\Filerobot\Model\Product\Gallery\UpdateHandlerOverride" />
+                   type="Scaleflex\Filerobot\Model\Product\Gallery\UpdateHandler" />
        <preference for="Magento\Catalog\Ui\Component\Listing\Columns\Thumbnail"
                    type="Scaleflex\Filerobot\Ui\Component\Listing\Columns\Thumbnail" />
+       <preference for="Magento\CatalogStaging\Model\Product\Gallery\CreateHandler"
+                   type="Scaleflex\Filerobot\Model\Product\Gallery\Staging\CreateHandler" />
+       <preference for="Magento\CatalogStaging\Model\Product\Gallery\UpdateHandler"
+                   type="Scaleflex\Filerobot\Model\Product\Gallery\Staging\UpdateHandler" />
 </config>
diff --git a/etc/module.xml b/etc/module.xml
index 24f8355..b2aa56d 100644
--- a/etc/module.xml
+++ b/etc/module.xml
@@ -1,8 +1,10 @@
 <?xml version="1.0" ?>
 <config xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:noNamespaceSchemaLocation="urn:magento:framework:Module/etc/module.xsd">
 	<module name="Scaleflex_Filerobot"  setup_version="1.0.0">
+	         <module name="Magento_Eav"/>
 	         <module name="Magento_Ui"/>
 	         <module name="Magento_Catalog"/>
+	         <module name="Magento_CatalogStaging"/>
 	         <module name="Magento_Customer"/>
              <module name="Magento_Store"/>
 	</module>
diff --git a/view/frontend/web/js/view/summary/item/details/thumbnail.js b/view/frontend/web/js/view/summary/item/details/thumbnail.js
index d78d439..6afc642 100644
--- a/view/frontend/web/js/view/summary/item/details/thumbnail.js
+++ b/view/frontend/web/js/view/summary/item/details/thumbnail.js
@@ -25,7 +25,7 @@ define(['uiComponent'], function (Component) {
 
         getQuoteItem: function (item) {
             const item_id = item['item_id'];
-            return this.quoteItemData.find((quoteItem) => quoteItem['item_id'] === item_id);
+            return this.quoteItemData.find((quoteItem) => String(quoteItem['item_id']) === String(item_id));
         },
 
         /**
@@ -34,8 +34,8 @@ define(['uiComponent'], function (Component) {
          */
         getSrc: function (item) {
             const quoteItem = this.getQuoteItem(item);
-            if (quoteItem['product']['thumbnail'].includes('filerobot')) {
-                return quoteItem['product']['thumbnail'];
+            if (quoteItem?.product?.thumbnail?.includes('filerobot')) {
+                return quoteItem?.product?.thumbnail;
             } else if (this.imageData[item['item_id']]) {
                 return this.imageData[item['item_id']].src;
             }