diff --git a/packages/grafana-flamegraph/src/FlameGraph/FlameGraph.test.tsx b/packages/grafana-flamegraph/src/FlameGraph/FlameGraph.test.tsx
index 70694baf986..fc152e96b00 100644
--- a/packages/grafana-flamegraph/src/FlameGraph/FlameGraph.test.tsx
+++ b/packages/grafana-flamegraph/src/FlameGraph/FlameGraph.test.tsx
@@ -1,7 +1,7 @@
 import { fireEvent, render, screen } from '@testing-library/react';
 import React from 'react';
 
-import { createDataFrame, createTheme } from '@grafana/data';
+import { createDataFrame } from '@grafana/data';
 
 import { ColorScheme } from '../types';
 
@@ -48,7 +48,6 @@ describe('FlameGraph', () => {
         onFocusPillClick={onFocusPillClick}
         onSandwichPillClick={onSandwichPillClick}
         colorScheme={ColorScheme.ValueBased}
-        getTheme={() => createTheme({ colors: { mode: 'dark' } })}
       />
     );
     return {
diff --git a/packages/grafana-flamegraph/src/FlameGraph/FlameGraph.tsx b/packages/grafana-flamegraph/src/FlameGraph/FlameGraph.tsx
index 0af2073284b..485d8b5e19a 100644
--- a/packages/grafana-flamegraph/src/FlameGraph/FlameGraph.tsx
+++ b/packages/grafana-flamegraph/src/FlameGraph/FlameGraph.tsx
@@ -20,7 +20,6 @@ import { css } from '@emotion/css';
 import React, { MouseEvent as ReactMouseEvent, useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import { useMeasure } from 'react-use';
 
-import { GrafanaTheme2 } from '@grafana/data';
 import { Icon } from '@grafana/ui';
 
 import { PIXELS_PER_LEVEL } from '../constants';
@@ -48,7 +47,6 @@ type Props = {
   onFocusPillClick: () => void;
   onSandwichPillClick: () => void;
   colorScheme: ColorScheme | ColorSchemeDiff;
-  getTheme: () => GrafanaTheme2;
 };
 
 const FlameGraph = ({
@@ -66,7 +64,6 @@ const FlameGraph = ({
   onFocusPillClick,
   onSandwichPillClick,
   colorScheme,
-  getTheme,
 }: Props) => {
   const styles = getStyles();
 
@@ -108,7 +105,6 @@ const FlameGraph = ({
     totalColorTicks: data.isDiffFlamegraph() ? totalProfileTicks : totalViewTicks,
     totalTicksRight: totalProfileTicksRight,
     wrapperWidth,
-    getTheme,
   });
 
   const onGraphClick = useCallback(
@@ -186,7 +182,6 @@ const FlameGraph = ({
   return (
     <div className={styles.graph}>
       <FlameGraphMetadata
-        getTheme={getTheme}
         data={data}
         focusedItem={focusedItemData}
         sandwichedLabel={sandwichItem}
@@ -220,13 +215,7 @@ const FlameGraph = ({
           />
         </div>
       </div>
-      <FlameGraphTooltip
-        getTheme={getTheme}
-        position={mousePosition}
-        item={tooltipItem}
-        data={data}
-        totalTicks={totalViewTicks}
-      />
+      <FlameGraphTooltip position={mousePosition} item={tooltipItem} data={data} totalTicks={totalViewTicks} />
       {clickedItemData && (
         <FlameGraphContextMenu
           itemData={clickedItemData}
diff --git a/packages/grafana-flamegraph/src/FlameGraph/FlameGraphMetadata.test.tsx b/packages/grafana-flamegraph/src/FlameGraph/FlameGraphMetadata.test.tsx
index feb3bf50a6f..a92270d6b7c 100644
--- a/packages/grafana-flamegraph/src/FlameGraph/FlameGraphMetadata.test.tsx
+++ b/packages/grafana-flamegraph/src/FlameGraph/FlameGraphMetadata.test.tsx
@@ -2,8 +2,6 @@ import { render, screen } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import React from 'react';
 
-import { createTheme } from '@grafana/data';
-
 import FlameGraphMetadata from './FlameGraphMetadata';
 import { textToDataContainer } from './testHelpers';
 
@@ -23,7 +21,6 @@ function setup(props: Partial<React.ComponentProps<typeof FlameGraphMetadata>> =
       totalTicks={17}
       onFocusPillClick={onFocusPillClick}
       onSandwichPillClick={onSandwichPillClick}
-      getTheme={() => createTheme({ colors: { mode: 'dark' } })}
       {...props}
     />
   );
diff --git a/packages/grafana-flamegraph/src/FlameGraph/FlameGraphMetadata.tsx b/packages/grafana-flamegraph/src/FlameGraph/FlameGraphMetadata.tsx
index 78073c5c90d..9a896f33a0d 100644
--- a/packages/grafana-flamegraph/src/FlameGraph/FlameGraphMetadata.tsx
+++ b/packages/grafana-flamegraph/src/FlameGraph/FlameGraphMetadata.tsx
@@ -2,7 +2,7 @@ import { css } from '@emotion/css';
 import React, { ReactNode } from 'react';
 
 import { getValueFormat, GrafanaTheme2 } from '@grafana/data';
-import { Icon, IconButton } from '@grafana/ui';
+import { Icon, IconButton, useStyles2 } from '@grafana/ui';
 
 import { ClickedItemData } from '../types';
 
@@ -15,12 +15,11 @@ type Props = {
   onSandwichPillClick: () => void;
   focusedItem?: ClickedItemData;
   sandwichedLabel?: string;
-  getTheme: () => GrafanaTheme2;
 };
 
 const FlameGraphMetadata = React.memo(
-  ({ data, focusedItem, totalTicks, sandwichedLabel, onFocusPillClick, onSandwichPillClick, getTheme }: Props) => {
-    const styles = getStyles(getTheme());
+  ({ data, focusedItem, totalTicks, sandwichedLabel, onFocusPillClick, onSandwichPillClick }: Props) => {
+    const styles = useStyles2(getStyles);
     const parts: ReactNode[] = [];
     const ticksVal = getValueFormat('short')(totalTicks);
 
diff --git a/packages/grafana-flamegraph/src/FlameGraph/FlameGraphTooltip.tsx b/packages/grafana-flamegraph/src/FlameGraph/FlameGraphTooltip.tsx
index 1949d2ba08c..fe9d7dd1316 100644
--- a/packages/grafana-flamegraph/src/FlameGraph/FlameGraphTooltip.tsx
+++ b/packages/grafana-flamegraph/src/FlameGraph/FlameGraphTooltip.tsx
@@ -2,7 +2,7 @@ import { css } from '@emotion/css';
 import React from 'react';
 
 import { DisplayValue, getValueFormat, GrafanaTheme2 } from '@grafana/data';
-import { InteractiveTable, Portal, VizTooltipContainer } from '@grafana/ui';
+import { InteractiveTable, Portal, useStyles2, VizTooltipContainer } from '@grafana/ui';
 
 import { FlameGraphDataContainer, LevelItem } from './dataTransform';
 
@@ -11,11 +11,10 @@ type Props = {
   totalTicks: number;
   position?: { x: number; y: number };
   item?: LevelItem;
-  getTheme: () => GrafanaTheme2;
 };
 
-const FlameGraphTooltip = ({ data, item, totalTicks, position, getTheme }: Props) => {
-  const styles = getStyles(getTheme());
+const FlameGraphTooltip = ({ data, item, totalTicks, position }: Props) => {
+  const styles = useStyles2(getStyles);
 
   if (!(item && position)) {
     return null;
diff --git a/packages/grafana-flamegraph/src/FlameGraph/rendering.ts b/packages/grafana-flamegraph/src/FlameGraph/rendering.ts
index b017ddbed2b..b38ade060a5 100644
--- a/packages/grafana-flamegraph/src/FlameGraph/rendering.ts
+++ b/packages/grafana-flamegraph/src/FlameGraph/rendering.ts
@@ -3,6 +3,7 @@ import { RefObject, useEffect, useMemo, useState } from 'react';
 import color from 'tinycolor2';
 
 import { GrafanaTheme2 } from '@grafana/data';
+import { useTheme2 } from '@grafana/ui';
 
 import {
   BAR_BORDER_WIDTH,
@@ -41,7 +42,6 @@ type RenderOptions = {
   totalTicksRight: number | undefined;
   colorScheme: ColorScheme | ColorSchemeDiff;
   focusedItemData?: ClickedItemData;
-  getTheme: () => GrafanaTheme2;
 };
 
 export function useFlameRender(options: RenderOptions) {
@@ -59,7 +59,6 @@ export function useFlameRender(options: RenderOptions) {
     totalTicksRight,
     colorScheme,
     focusedItemData,
-    getTheme,
   } = options;
   const foundLabels = useMemo(() => {
     if (search) {
@@ -79,7 +78,7 @@ export function useFlameRender(options: RenderOptions) {
   }, [search, data]);
 
   const ctx = useSetupCanvas(canvasRef, wrapperWidth, levels.length);
-  const theme = getTheme();
+  const theme = useTheme2();
 
   useEffect(() => {
     if (!ctx) {
diff --git a/packages/grafana-flamegraph/src/FlameGraphContainer.tsx b/packages/grafana-flamegraph/src/FlameGraphContainer.tsx
index 18970d094b2..6aacd816674 100644
--- a/packages/grafana-flamegraph/src/FlameGraphContainer.tsx
+++ b/packages/grafana-flamegraph/src/FlameGraphContainer.tsx
@@ -3,6 +3,7 @@ import React, { useCallback, useEffect, useMemo, useState } from 'react';
 import { useMeasure } from 'react-use';
 
 import { DataFrame, GrafanaTheme2 } from '@grafana/data';
+import { ThemeContext } from '@grafana/ui';
 
 import FlameGraph from './FlameGraph/FlameGraph';
 import { FlameGraphDataContainer } from './FlameGraph/dataTransform';
@@ -133,75 +134,76 @@ const FlameGraphContainer = ({
   }
 
   return (
-    <div ref={sizeRef} className={styles.container}>
-      <FlameGraphHeader
-        search={search}
-        setSearch={setSearch}
-        selectedView={selectedView}
-        setSelectedView={(view) => {
-          setSelectedView(view);
-          onViewSelected?.(view);
-        }}
-        containerWidth={containerWidth}
-        onReset={() => {
-          resetFocus();
-          resetSandwich();
-        }}
-        textAlign={textAlign}
-        onTextAlignChange={(align) => {
-          setTextAlign(align);
-          onTextAlignSelected?.(align);
-        }}
-        showResetButton={Boolean(focusedItemData || sandwichItem)}
-        colorScheme={colorScheme}
-        onColorSchemeChange={setColorScheme}
-        stickyHeader={Boolean(stickyHeader)}
-        extraHeaderElements={extraHeaderElements}
-        vertical={vertical}
-        isDiffMode={Boolean(dataContainer.isDiffFlamegraph())}
-        getTheme={getTheme}
-      />
+    // We add the theme context to bridge the gap if this is rendered in non grafana environment where the context
+    // isn't already provided.
+    <ThemeContext.Provider value={theme}>
+      <div ref={sizeRef} className={styles.container}>
+        <FlameGraphHeader
+          search={search}
+          setSearch={setSearch}
+          selectedView={selectedView}
+          setSelectedView={(view) => {
+            setSelectedView(view);
+            onViewSelected?.(view);
+          }}
+          containerWidth={containerWidth}
+          onReset={() => {
+            resetFocus();
+            resetSandwich();
+          }}
+          textAlign={textAlign}
+          onTextAlignChange={(align) => {
+            setTextAlign(align);
+            onTextAlignSelected?.(align);
+          }}
+          showResetButton={Boolean(focusedItemData || sandwichItem)}
+          colorScheme={colorScheme}
+          onColorSchemeChange={setColorScheme}
+          stickyHeader={Boolean(stickyHeader)}
+          extraHeaderElements={extraHeaderElements}
+          vertical={vertical}
+          isDiffMode={Boolean(dataContainer.isDiffFlamegraph())}
+        />
 
-      <div className={styles.body}>
-        {selectedView !== SelectedView.FlameGraph && (
-          <FlameGraphTopTableContainer
-            data={dataContainer}
-            onSymbolClick={onSymbolClick}
-            height={selectedView === SelectedView.TopTable ? 600 : undefined}
-            search={search}
-            sandwichItem={sandwichItem}
-            onSandwich={setSandwichItem}
-            onSearch={setSearch}
-            onTableSort={onTableSort}
-            getTheme={getTheme}
-            vertical={vertical}
-          />
-        )}
+        <div className={styles.body}>
+          {selectedView !== SelectedView.FlameGraph && (
+            <FlameGraphTopTableContainer
+              data={dataContainer}
+              onSymbolClick={onSymbolClick}
+              height={selectedView === SelectedView.TopTable ? 600 : undefined}
+              search={search}
+              sandwichItem={sandwichItem}
+              onSandwich={setSandwichItem}
+              onSearch={setSearch}
+              onTableSort={onTableSort}
+              vertical={vertical}
+            />
+          )}
 
-        {selectedView !== SelectedView.TopTable && (
-          <FlameGraph
-            getTheme={getTheme}
-            data={dataContainer}
-            rangeMin={rangeMin}
-            rangeMax={rangeMax}
-            search={search}
-            setRangeMin={setRangeMin}
-            setRangeMax={setRangeMax}
-            onItemFocused={(data) => setFocusedItemData(data)}
-            focusedItemData={focusedItemData}
-            textAlign={textAlign}
-            sandwichItem={sandwichItem}
-            onSandwich={(label: string) => {
-              resetFocus();
-              setSandwichItem(label);
-            }}
-            onFocusPillClick={resetFocus}
-            onSandwichPillClick={resetSandwich}
-            colorScheme={colorScheme}
-          />
-        )}
+          {selectedView !== SelectedView.TopTable && (
+            <FlameGraph
+              data={dataContainer}
+              rangeMin={rangeMin}
+              rangeMax={rangeMax}
+              search={search}
+              setRangeMin={setRangeMin}
+              setRangeMax={setRangeMax}
+              onItemFocused={(data) => setFocusedItemData(data)}
+              focusedItemData={focusedItemData}
+              textAlign={textAlign}
+              sandwichItem={sandwichItem}
+              onSandwich={(label: string) => {
+                resetFocus();
+                setSandwichItem(label);
+              }}
+              onFocusPillClick={resetFocus}
+              onSandwichPillClick={resetSandwich}
+              colorScheme={colorScheme}
+            />
+          )}
+        </div>
       </div>
-    </div>
+    </ThemeContext.Provider>
   );
 };
 
diff --git a/packages/grafana-flamegraph/src/FlameGraphHeader.test.tsx b/packages/grafana-flamegraph/src/FlameGraphHeader.test.tsx
index edc0f9ddb42..76935d7d198 100644
--- a/packages/grafana-flamegraph/src/FlameGraphHeader.test.tsx
+++ b/packages/grafana-flamegraph/src/FlameGraphHeader.test.tsx
@@ -3,8 +3,6 @@ import { render, screen } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import React from 'react';
 
-import { createTheme } from '@grafana/data';
-
 import FlameGraphHeader from './FlameGraphHeader';
 import { ColorScheme, SelectedView } from './types';
 
@@ -28,7 +26,6 @@ describe('FlameGraphHeader', () => {
         showResetButton={true}
         colorScheme={ColorScheme.ValueBased}
         onColorSchemeChange={onSchemeChange}
-        getTheme={() => createTheme({ colors: { mode: 'dark' } })}
         stickyHeader={false}
         isDiffMode={false}
         {...props}
diff --git a/packages/grafana-flamegraph/src/FlameGraphHeader.tsx b/packages/grafana-flamegraph/src/FlameGraphHeader.tsx
index 9d3f601d8d8..7c521af4714 100644
--- a/packages/grafana-flamegraph/src/FlameGraphHeader.tsx
+++ b/packages/grafana-flamegraph/src/FlameGraphHeader.tsx
@@ -4,7 +4,7 @@ import useDebounce from 'react-use/lib/useDebounce';
 import usePrevious from 'react-use/lib/usePrevious';
 
 import { GrafanaTheme2, SelectableValue } from '@grafana/data';
-import { Button, Dropdown, Input, Menu, RadioButtonGroup } from '@grafana/ui';
+import { Button, Dropdown, Input, Menu, RadioButtonGroup, useStyles2 } from '@grafana/ui';
 
 import { byPackageGradient, byValueGradient, diffColorBlindGradient, diffDefaultGradient } from './FlameGraph/colors';
 import { MIN_WIDTH_TO_SHOW_BOTH_TOPTABLE_AND_FLAMEGRAPH } from './constants';
@@ -25,7 +25,6 @@ type Props = {
   stickyHeader: boolean;
   vertical?: boolean;
   isDiffMode: boolean;
-  getTheme: () => GrafanaTheme2;
 
   extraHeaderElements?: React.ReactNode;
 };
@@ -46,9 +45,8 @@ const FlameGraphHeader = ({
   extraHeaderElements,
   vertical,
   isDiffMode,
-  getTheme,
 }: Props) => {
-  const styles = getStyles(getTheme(), stickyHeader);
+  const styles = useStyles2(getStyles, stickyHeader);
   const [localSearch, setLocalSearch] = useSearchInput(search, setSearch);
 
   const suffix =
@@ -95,12 +93,7 @@ const FlameGraphHeader = ({
             aria-label={'Reset focus and sandwich state'}
           />
         )}
-        <ColorSchemeButton
-          value={colorScheme}
-          onChange={onColorSchemeChange}
-          isDiffMode={isDiffMode}
-          getTheme={getTheme}
-        />
+        <ColorSchemeButton value={colorScheme} onChange={onColorSchemeChange} isDiffMode={isDiffMode} />
         <RadioButtonGroup<TextAlign>
           size="sm"
           disabled={selectedView === SelectedView.TopTable}
@@ -124,12 +117,11 @@ const FlameGraphHeader = ({
 type ColorSchemeButtonProps = {
   value: ColorScheme | ColorSchemeDiff;
   onChange: (colorScheme: ColorScheme | ColorSchemeDiff) => void;
-  getTheme: () => GrafanaTheme2;
   isDiffMode: boolean;
 };
 function ColorSchemeButton(props: ColorSchemeButtonProps) {
   // TODO: probably create separate getStyles
-  const styles = getStyles(props.getTheme(), false);
+  const styles = useStyles2(getStyles, false);
   let menu = (
     <Menu>
       <Menu.Item label="By package name" onClick={() => props.onChange(ColorScheme.PackageBased)} />
diff --git a/packages/grafana-flamegraph/src/TopTable/FlameGraphTopTableContainer.test.tsx b/packages/grafana-flamegraph/src/TopTable/FlameGraphTopTableContainer.test.tsx
index be5f85a64ec..be543a09f95 100644
--- a/packages/grafana-flamegraph/src/TopTable/FlameGraphTopTableContainer.test.tsx
+++ b/packages/grafana-flamegraph/src/TopTable/FlameGraphTopTableContainer.test.tsx
@@ -2,7 +2,7 @@ import { render, screen } from '@testing-library/react';
 import userEvents from '@testing-library/user-event';
 import React from 'react';
 
-import { createDataFrame, createTheme } from '@grafana/data';
+import { createDataFrame } from '@grafana/data';
 
 import { FlameGraphDataContainer } from '../FlameGraph/dataTransform';
 import { data } from '../FlameGraph/testData/dataNestedSet';
@@ -22,7 +22,6 @@ describe('FlameGraphTopTableContainer', () => {
         onSymbolClick={jest.fn()}
         onSearch={onSearch}
         onSandwich={onSandwich}
-        getTheme={() => createTheme({ colors: { mode: 'dark' } })}
       />
     );
 
diff --git a/packages/grafana-flamegraph/src/TopTable/FlameGraphTopTableContainer.tsx b/packages/grafana-flamegraph/src/TopTable/FlameGraphTopTableContainer.tsx
index 51de2965308..cef80bc9806 100644
--- a/packages/grafana-flamegraph/src/TopTable/FlameGraphTopTableContainer.tsx
+++ b/packages/grafana-flamegraph/src/TopTable/FlameGraphTopTableContainer.tsx
@@ -18,6 +18,8 @@ import {
   TableCustomCellOptions,
   TableFieldOptions,
   TableSortByFieldState,
+  useStyles2,
+  useTheme2,
 } from '@grafana/ui';
 
 import { FlameGraphDataContainer } from '../FlameGraph/dataTransform';
@@ -33,23 +35,11 @@ type Props = {
   onSearch: (str: string) => void;
   onSandwich: (str?: string) => void;
   onTableSort?: (sort: string) => void;
-  getTheme: () => GrafanaTheme2;
   vertical?: boolean;
 };
 
 const FlameGraphTopTableContainer = React.memo(
-  ({
-    data,
-    onSymbolClick,
-    height,
-    search,
-    onSearch,
-    sandwichItem,
-    onSandwich,
-    onTableSort,
-    getTheme,
-    vertical,
-  }: Props) => {
+  ({ data, onSymbolClick, height, search, onSearch, sandwichItem, onSandwich, onTableSort, vertical }: Props) => {
     const table = useMemo(() => {
       // Group the data by label, we show only one row per label and sum the values
       // TODO: should be by filename + funcName + linenumber?
@@ -73,7 +63,9 @@ const FlameGraphTopTableContainer = React.memo(
     // so we don't show potentially thousands of rows at once which can hinder performance (the table is virtualized
     // so with some max height it handles it fine)
     const tableHeight = vertical ? Math.min(Object.keys(table).length * rowHeight, 800) : 0;
-    const styles = getStyles(tableHeight, getTheme());
+
+    const styles = useStyles2(getStyles, tableHeight);
+    const theme = useTheme2();
 
     const [sort, setSort] = useState<TableSortByFieldState[]>([{ displayName: 'Self', desc: true }]);
 
@@ -92,7 +84,7 @@ const FlameGraphTopTableContainer = React.memo(
               onSymbolClick,
               onSearch,
               onSandwich,
-              getTheme,
+              theme,
               search,
               sandwichItem
             );
@@ -126,7 +118,7 @@ function buildTableDataFrame(
   onSymbolClick: (str: string) => void,
   onSearch: (str: string) => void,
   onSandwich: (str?: string) => void,
-  getTheme: () => GrafanaTheme2,
+  theme: GrafanaTheme2,
   search?: string,
   sandwichItem?: string
 ): DataFrame {
@@ -218,7 +210,7 @@ function buildTableDataFrame(
       overrides: [],
     },
     replaceVariables: (value: string) => value,
-    theme: getTheme(),
+    theme,
   });
 
   return dataFrames[0];
@@ -327,7 +319,7 @@ function ActionCell(props: ActionCellProps) {
   );
 }
 
-const getStyles = (height: number, theme: GrafanaTheme2) => {
+const getStyles = (theme: GrafanaTheme2, height: number) => {
   return {
     topTableContainer: css`
       label: topTableContainer;