diff --git a/packages/grafana-flamegraph/src/FlameGraphContainer.tsx b/packages/grafana-flamegraph/src/FlameGraphContainer.tsx
index bee6c307839..18970d094b2 100644
--- a/packages/grafana-flamegraph/src/FlameGraphContainer.tsx
+++ b/packages/grafana-flamegraph/src/FlameGraphContainer.tsx
@@ -206,24 +206,13 @@ const FlameGraphContainer = ({
 };
 
 function useColorScheme(dataContainer: FlameGraphDataContainer | undefined) {
-  const [colorScheme, setColorScheme] = useState<ColorScheme | ColorSchemeDiff>(
-    dataContainer?.isDiffFlamegraph() ? ColorSchemeDiff.Default : ColorScheme.ValueBased
-  );
-  useEffect(() => {
-    if (
-      dataContainer?.isDiffFlamegraph() &&
-      (colorScheme === ColorScheme.ValueBased || colorScheme === ColorScheme.PackageBased)
-    ) {
-      setColorScheme(ColorSchemeDiff.Default);
-    }
+  const defaultColorScheme = dataContainer?.isDiffFlamegraph() ? ColorSchemeDiff.Default : ColorScheme.PackageBased;
+  const [colorScheme, setColorScheme] = useState<ColorScheme | ColorSchemeDiff>(defaultColorScheme);
 
-    if (
-      !dataContainer?.isDiffFlamegraph() &&
-      (colorScheme === ColorSchemeDiff.Default || colorScheme === ColorSchemeDiff.DiffColorBlind)
-    ) {
-      setColorScheme(ColorScheme.ValueBased);
-    }
-  }, [dataContainer, colorScheme]);
+  // This makes sure that if we change the data to/from diff profile we reset the color scheme.
+  useEffect(() => {
+    setColorScheme(defaultColorScheme);
+  }, [defaultColorScheme]);
 
   return [colorScheme, setColorScheme] as const;
 }
diff --git a/packages/grafana-flamegraph/src/FlameGraphHeader.tsx b/packages/grafana-flamegraph/src/FlameGraphHeader.tsx
index 178efeb4b3e..9d3f601d8d8 100644
--- a/packages/grafana-flamegraph/src/FlameGraphHeader.tsx
+++ b/packages/grafana-flamegraph/src/FlameGraphHeader.tsx
@@ -132,8 +132,8 @@ function ColorSchemeButton(props: ColorSchemeButtonProps) {
   const styles = getStyles(props.getTheme(), false);
   let menu = (
     <Menu>
-      <Menu.Item label="By value" onClick={() => props.onChange(ColorScheme.ValueBased)} />
       <Menu.Item label="By package name" onClick={() => props.onChange(ColorScheme.PackageBased)} />
+      <Menu.Item label="By value" onClick={() => props.onChange(ColorScheme.ValueBased)} />
     </Menu>
   );