diff --git a/apps/remix-ide/src/app/plugins/solidity-umlgen.tsx b/apps/remix-ide/src/app/plugins/solidity-umlgen.tsx
index 72ee6a4100..6d9e4ee830 100644
--- a/apps/remix-ide/src/app/plugins/solidity-umlgen.tsx
+++ b/apps/remix-ide/src/app/plugins/solidity-umlgen.tsx
@@ -72,7 +72,7 @@ export class SolidityUmlGen extends ViewPlugin implements ISolidityUmlGen {
       const currentTheme: ThemeQualityType = await this.call('theme', 'currentTheme')
       this.currentlySelectedTheme = currentTheme.quality
       this.themeName = currentTheme.name
-      let result = ''
+      let result
       try {
         if (data.sources && Object.keys(data.sources).length > 1) { // we should flatten first as there are multiple asts
           result = await this.flattenContract(source, this.currentFile, data)
@@ -134,12 +134,12 @@ export class SolidityUmlGen extends ViewPlugin implements ISolidityUmlGen {
    */
   async flattenContract (source: any, filePath: string, data: any) {
     const hold = { data, source, filePath }
-    const ast = data.sources
-    const dependencyGraph = getDependencyGraph(ast, filePath)
+    const ast = hold.data.sources
+    const dependencyGraph = getDependencyGraph(ast, hold.filePath)
     const sorted = dependencyGraph.isEmpty()
         ? [filePath]
         : dependencyGraph.sort().reverse()
-    const sources = source.sources
+    const sources = hold.source.sources
     const result = concatSourceFiles(sorted, sources)
     await this.call('fileManager', 'writeFile', `${filePath}_flattened.sol`, result)
     return result
diff --git a/libs/remix-ui/solidity-uml-gen/src/lib/solidity-uml-gen.tsx b/libs/remix-ui/solidity-uml-gen/src/lib/solidity-uml-gen.tsx
index 42a1c315a6..6b679808da 100644
--- a/libs/remix-ui/solidity-uml-gen/src/lib/solidity-uml-gen.tsx
+++ b/libs/remix-ui/solidity-uml-gen/src/lib/solidity-uml-gen.tsx
@@ -40,19 +40,6 @@ export function RemixUiSolidityUmlGen ({ updatedSvg, loading, plugin }: RemixUiS
   const final = btoa(String.fromCharCode.apply(null, data))
 
   function ActionButtons({ actions: { zoomIn, zoomOut, resetTransform }}: ActionButtonsProps) {
-    const [reloadButton, setReloadButton] = useState(false)
-
-    useEffect(() => {
-      plugin.on('theme', 'themeChanged', (themeName) => {
-        console.log('theme changed', themeName)
-        setReloadButton(true)
-      })
-
-      return () => {
-        plugin.off('theme', 'themeChanged')
-        console.log('unmounted')
-      }
-    }, [])
   
     return (
       <>
@@ -69,19 +56,19 @@ export function RemixUiSolidityUmlGen ({ updatedSvg, loading, plugin }: RemixUiS
               <i className="far fa-arrow-to-bottom align-item-center d-flex justify-content-center"></i>
             </button>
             <button
-              className="btn btn-outline-info align-item-center d-flex justify-content-center rounded-circle mr-2"
+              className="badge badge-info p-2 align-item-center d-flex justify-content-center rounded-circle mr-2"
               onClick={() => zoomIn()}
             >
               <i className="far fa-plus "></i>
             </button>
             <button
-              className="btn btn-outline-info rounded-circle mr-2"
+              className="badge badge-info p-2 rounded-circle mr-2"
               onClick={() => zoomOut()}
             >
               <i className="far fa-minus align-item-center d-flex justify-content-center"></i>
             </button>
             <button
-              className="btn btn-outline-info rounded-circle mr-2"
+              className="badge badge-info p-2 rounded-circle mr-2"
               onClick={() => resetTransform()}
             >
               <i className="far fa-undo align-item-center d-flex justify-content-center"></i>