2 files changed, 20 insertions, 3 deletions

diff --git a/visualizer.c b/visualizer.c
index ff99e960f..bbb00debc 100644
--- a/visualizer.c
+++ b/visualizer.c
@@ -85,6 +85,15 @@ static remote_object_t* remote_objects[] = {
 GDisplay* LCD_DISPLAY = 0;
 GDisplay* LED_DISPLAY = 0;
 
+__attribute__((weak))
+GDisplay* get_lcd_display(void) {
+    return gdispGetDisplay(0);
+}
+
+__attribute__((weak))
+GDisplay* get_led_display(void) {
+    return gdispGetDisplay(1);
+}
 
 void start_keyframe_animation(keyframe_animation_t* animation) {
     animation->current_frame = -1;
@@ -444,10 +453,13 @@ void visualizer_init(void) {
 #ifdef USE_SERIAL_LINK
     add_remote_objects(remote_objects, sizeof(remote_objects) / sizeof(remote_object_t*) );
 #endif
-    // TODO: Make sure these works when either of these are disabled
-    LCD_DISPLAY = gdispGetDisplay(0);
-    LED_DISPLAY = gdispGetDisplay(1);
 
+#ifdef LCD_ENABLE
+    LCD_DISPLAY = get_lcd_display();
+#endif
+#ifdef LED_ENABLE
+    LED_DISPLAY = get_led_display();
+#endif
 
     // We are using a low priority thread, the idea is to have it run only
     // when the main thread is sleeping during the matrix scanning
diff --git a/visualizer.h b/visualizer.h
index 8a2772c6d..4d6a61dda 100644
--- a/visualizer.h
+++ b/visualizer.h
@@ -45,6 +45,11 @@ void visualizer_suspend(void);
 // This should be called when the keyboard wakes up from suspend state
 void visualizer_resume(void);
 
+// These functions are week, so they can be overridden by the keyboard
+// if needed
+GDisplay* get_lcd_display(void);
+GDisplay* get_led_display(void);
+
 // If you need support for more than 16 keyframes per animation, you can change this
 #define MAX_VISUALIZER_KEY_FRAMES 16