diff --git a/src/ui/lib/text/HighlightGrammar.tsx b/src/ui/lib/text/HighlightGrammar.tsx
index 4186f9f..c567e6d 100644
--- a/src/ui/lib/text/HighlightGrammar.tsx
+++ b/src/ui/lib/text/HighlightGrammar.tsx
@@ -14,8 +14,6 @@ import bash from 'react-syntax-highlighter/dist/esm/languages/prism/bash';
 import dark from 'react-syntax-highlighter/dist/esm/styles/prism/material-dark';
 import light from 'react-syntax-highlighter/dist/esm/styles/prism/material-light';
 
-import useMode from '../../theme/useMode';
-import useIsMounted from '../component/useIsMounted';
 import saveTextToClipboard from '../output/saveTextToClipboard';
 
 SyntaxHighlighter.registerLanguage('python', python);
@@ -59,13 +57,21 @@ const copyToClipboardColor = (state: State) => {
 
 const customStyle = {margin: 0, display: 'flex', flex: '1'};
 
-type Style = typeof dark | typeof light;
+function prefixKeys<T>(
+	prefix: string,
+	record: Record<string, T>,
+): Record<string, T> {
+	return Object.fromEntries(
+		Object.entries(record).map(([k, v]) => [prefix + k, v]),
+	);
+}
 
-function HighlightGrammar({language, word}: Props) {
-	const isMounted = useIsMounted();
+const style = {
+	...prefixKeys('html[data-mui-color-scheme="dark"] ', dark),
+	...light,
+};
 
-	const mode = useMode();
-	const style: Style = mode === 'dark' ? dark : light;
+function HighlightGrammar({language, word}: Props) {
 	const [tooltipText, setTooltipText] = useState<State>(init);
 	const [open, setOpen] = useState<boolean>(false);
 	const handleOpen = () => {
@@ -81,11 +87,6 @@ function HighlightGrammar({language, word}: Props) {
 		setOpen(false);
 	}, [language, word]);
 
-	if (!isMounted()) {
-		// NOTE We cannot highlight with the correct mode on the server
-		return null;
-	}
-
 	return (
 		<Paper variant="outlined" sx={{display: 'flex', position: 'relative'}}>
 			<SyntaxHighlighter