diff --git a/components/jiecang_desk_controller/__init__.py b/components/jiecang_desk_controller/__init__.py
index 7e7e395..1415ee0 100644
--- a/components/jiecang_desk_controller/__init__.py
+++ b/components/jiecang_desk_controller/__init__.py
@@ -108,19 +108,19 @@ CONFIG_SCHEMA = cv.COMPONENT_SCHEMA.extend({
         }),
     }),
     cv.Optional(CONF_BUTTONS): cv.Schema({
-        cv.Optional(CONF_STEP_UP): button.BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
-        cv.Optional(CONF_STEP_DOWN): button.BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
-        cv.Optional(CONF_STOP): button.BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
-        cv.Optional(CONF_POSITION1): button.BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
-        cv.Optional(CONF_POSITION2): button.BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
-        cv.Optional(CONF_POSITION3): button.BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
-        cv.Optional(CONF_POSITION4): button.BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
-        cv.Optional(CONF_SAVE_POSITION): button.BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
-        cv.Optional(CONF_MOVE_UP): button.BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
-        cv.Optional(CONF_MOVE_DOWN): button.BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
+        cv.Optional(CONF_STEP_UP): button._BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
+        cv.Optional(CONF_STEP_DOWN): button._BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
+        cv.Optional(CONF_STOP): button._BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
+        cv.Optional(CONF_POSITION1): button._BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
+        cv.Optional(CONF_POSITION2): button._BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
+        cv.Optional(CONF_POSITION3): button._BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
+        cv.Optional(CONF_POSITION4): button._BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
+        cv.Optional(CONF_SAVE_POSITION): button._BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
+        cv.Optional(CONF_MOVE_UP): button._BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
+        cv.Optional(CONF_MOVE_DOWN): button._BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
         # Backward compatibility
-        cv.Optional(CONF_RAISE): button.BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
-        cv.Optional(CONF_LOWER): button.BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
+        cv.Optional(CONF_RAISE): button._BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
+        cv.Optional(CONF_LOWER): button._BUTTON_SCHEMA.extend({cv.GenerateID(): cv.declare_id(JiecangDeskButton)}),
     }),
 }).extend(uart.UART_DEVICE_SCHEMA)