diff --git a/packages/mermaid/src/diagrams/class/parser/antlr/ClassLexer.g4 b/packages/mermaid/src/diagrams/class/parser/antlr/ClassLexer.g4
index 8e51dc24b..442312126 100644
--- a/packages/mermaid/src/diagrams/class/parser/antlr/ClassLexer.g4
+++ b/packages/mermaid/src/diagrams/class/parser/antlr/ClassLexer.g4
@@ -1,4 +1,6 @@
 lexer grammar ClassLexer;
+import HeaderCommon;
+
 
 tokens {
   ACC_TITLE_VALUE,
@@ -28,13 +30,13 @@ fragment NOT_DQUOTE: ~[""];
 
 
 // Comments and whitespace
-COMMENT: '%%' ~[\r\n]* -> skip;
+
 NEWLINE: ('\r'? '\n')+ { this.clearPendingScopes(); };
 WS: [ \t]+ -> skip;
 
 // Diagram title declaration
-CLASS_DIAGRAM_V2: 'classDiagram-v2' -> type(CLASS_DIAGRAM);
-CLASS_DIAGRAM: 'classDiagram';
+CLASS_DIAGRAM_V2: 'classDiagram-v2' { this.headerMode = false; } -> type(CLASS_DIAGRAM);
+CLASS_DIAGRAM: 'classDiagram' { this.headerMode = false; };
 
 // Directions
 DIRECTION_TB: 'direction' WS_INLINE+ 'TB';
diff --git a/packages/mermaid/src/diagrams/flowchart/parser/antlr/FlowLexer.g4 b/packages/mermaid/src/diagrams/flowchart/parser/antlr/FlowLexer.g4
index 09220cee5..3b97ec8c5 100644
--- a/packages/mermaid/src/diagrams/flowchart/parser/antlr/FlowLexer.g4
+++ b/packages/mermaid/src/diagrams/flowchart/parser/antlr/FlowLexer.g4
@@ -1,4 +1,6 @@
 lexer grammar FlowLexer;
+import HeaderCommon;
+
 
 // Virtual tokens for parser
 tokens {
@@ -25,7 +27,7 @@ HREF: 'href' WS;
 CLICK: 'click' WS+ [A-Za-z0-9_]+ -> pushMode(CLICK_MODE);
 
 // Graph declaration tokens - these trigger direction mode
-GRAPH: ('flowchart-elk' | 'graph' | 'flowchart') -> pushMode(DIR_MODE);
+GRAPH: ('flowchart-elk' | 'graph' | 'flowchart') { this.headerMode = false; } -> pushMode(DIR_MODE);
 SUBGRAPH: 'subgraph';
 END: 'end';