diff --git a/components/state_machine/__init__.py b/components/state_machine/__init__.py
index 29f5828..c8d7f26 100644
--- a/components/state_machine/__init__.py
+++ b/components/state_machine/__init__.py
@@ -94,7 +94,7 @@ def validate_transition(value):
     return validate_transition({CONF_FROM: a, CONF_TO: b})
 
 def output_graph(config):
-    if config[CONF_DIAGRAM] == None or config[CONF_DIAGRAM] == False:
+    if not CONF_DIAGRAM in config:
         return config
 
     if config[CONF_DIAGRAM] == "mermaid":