diff --git a/icnss2/main.h b/icnss2/main.h
index 27c4ef8161..d96c536dd5 100644
--- a/icnss2/main.h
+++ b/icnss2/main.h
@@ -10,7 +10,11 @@
 #include <linux/kobject.h>
 #include <linux/platform_device.h>
 #include <linux/ipc_logging.h>
+#ifdef CONFIG_CNSS_OUT_OF_TREE
+#include "icnss2.h"
+#else
 #include <soc/qcom/icnss2.h>
+#endif
 #include "wlan_firmware_service_v01.h"
 #include <linux/mailbox_client.h>
 
diff --git a/icnss2/qmi.c b/icnss2/qmi.c
index 478c03f6ce..f9f37aaf79 100644
--- a/icnss2/qmi.c
+++ b/icnss2/qmi.c
@@ -22,9 +22,11 @@
 #include <linux/firmware.h>
 #include <linux/soc/qcom/qmi.h>
 #include <linux/platform_device.h>
+#ifdef CONFIG_CNSS_OUT_OF_TREE
+#include "icnss2.h"
+#else
 #include <soc/qcom/icnss2.h>
-#include <soc/qcom/service-locator.h>
-#include <soc/qcom/service-notifier.h>
+#endif
 #include <soc/qcom/of_common.h>
 #include "wlan_firmware_service_v01.h"
 #include "main.h"