From 6b8c1dde93396a32506f4eede244fd4e1aa689a3 Mon Sep 17 00:00:00 2001
From: rcmerci <rcmerci@gmail.com>
Date: Tue, 18 Mar 2025 15:43:44 +0800
Subject: [PATCH] feat(embedding): check webgpu available

---
 src/main/frontend/handler.cljs            |  9 ++++++---
 src/main/frontend/persist_db/browser.cljs | 10 ++++++++--
 src/main/frontend/worker/embedding.cljs   | 10 ++++------
 3 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/src/main/frontend/handler.cljs b/src/main/frontend/handler.cljs
index 83280d3428..66ecb6efd7 100644
--- a/src/main/frontend/handler.cljs
+++ b/src/main/frontend/handler.cljs
@@ -169,15 +169,18 @@
    (when (mobile-util/native-platform?)
      (mobile/mobile-preinit))
    (-> (p/let [_ (db-browser/start-db-worker!)
-               _ (db-browser/start-inference-worker!)
-               _ (db-browser/<connect-db-worker-and-infer-worker!)
                repos (repo-handler/get-repos)
                _ (state/set-repos! repos)
                _ (mobile-util/hide-splash) ;; hide splash as early as ui is stable
                repo (or (state/get-current-repo) (:url (first repos)))
                _ (if (empty? repos)
                    (repo-handler/new-db! config/demo-repo)
-                   (restore-and-setup! repo))]
+                   (restore-and-setup! repo))
+               webgpu-available? (db-browser/<check-webgpu-available?)]
+         (log/info :webgpu-available? webgpu-available?)
+         (when webgpu-available?
+           (p/do! (db-browser/start-inference-worker!)
+                  (db-browser/<connect-db-worker-and-infer-worker!)))
          (when (util/electron?)
            (persist-db/run-export-periodically!))
          (when (mobile-util/native-platform?)
diff --git a/src/main/frontend/persist_db/browser.cljs b/src/main/frontend/persist_db/browser.cljs
index 66c825e7bc..07184d1638 100644
--- a/src/main/frontend/persist_db/browser.cljs
+++ b/src/main/frontend/persist_db/browser.cljs
@@ -149,6 +149,12 @@
                      (js/console.error error)
                      (notification/show! "It seems that OPFS is not supported on this browser, please upgrade this browser to the latest version or use another browser." :error)))))))
 
+(defn <check-webgpu-available?
+  []
+  (if (some? js/navigator.gpu)
+    (p/chain (js/navigator.gpu.requestAdapter) some?)
+    (p/promise false)))
+
 (defn start-inference-worker!
   []
   (when-not util/node-test?
@@ -167,8 +173,8 @@
   []
   (assert (and @state/*infer-worker @state/*db-worker))
   (p/do!
-    (.set-db-worker-proxy ^js @state/*infer-worker (Comlink/proxy @state/*db-worker))
-    (.set-infer-worker-proxy ^js @state/*db-worker (Comlink/proxy @state/*infer-worker))))
+   (.set-db-worker-proxy ^js @state/*infer-worker (Comlink/proxy @state/*db-worker))
+   (.set-infer-worker-proxy ^js @state/*db-worker (Comlink/proxy @state/*infer-worker))))
 
 (defn <export-db!
   [repo data]
diff --git a/src/main/frontend/worker/embedding.cljs b/src/main/frontend/worker/embedding.cljs
index 61161a99ff..51948b7b7f 100644
--- a/src/main/frontend/worker/embedding.cljs
+++ b/src/main/frontend/worker/embedding.cljs
@@ -12,6 +12,7 @@
 ;;; TODOs:
 ;;; - [x] add :logseq.property/description into text-to-embedding
 ;;; - add tags to text-to-embedding
+;;; - [x] check webgpu available, transformers.js is slow without webgpu(the difference is ~70 times)
 
 (defn- stale-block-filter-preds
   "When `reset?`, ignore :logseq.property.embedding/hnsw-label-updated-at in block"
@@ -81,8 +82,7 @@
   outdate rule: block/updated-at > :logseq.property.embedding/hnsw-label-updated-at"
   [repo conn]
   (m/sp
-    (let [^js infer-worker @worker-state/*infer-worker]
-      (assert (some? infer-worker))
+    (when-let [^js infer-worker @worker-state/*infer-worker]
       (let [stale-blocks (stale-block-lazy-seq @conn false)]
         (doseq [stale-block-chunk (sequence (partition-by-text-size 2000) stale-blocks)]
           (let [e+updated-at-coll (map (juxt :db/id :block/updated-at) stale-block-chunk)
@@ -100,8 +100,7 @@
   "force re-embedding all block-data in graph"
   [repo conn]
   (m/sp
-    (let [^js infer-worker @worker-state/*infer-worker]
-      (assert (some? infer-worker))
+    (when-let [^js infer-worker @worker-state/*infer-worker]
       (c.m/<? (.force-reset-index! infer-worker repo))
       (let [all-blocks (stale-block-lazy-seq @conn true)]
         (doseq [block-chunk (sequence (partition-by-text-size 2000) all-blocks)]
@@ -128,8 +127,7 @@
 (defn <search
   [repo conn query-string nums-neighbors]
   (m/sp
-    (let [^js infer-worker @worker-state/*infer-worker]
-      (assert (some? infer-worker))
+    (when-let [^js infer-worker @worker-state/*infer-worker]
       (let [{:keys [distances neighbors] :as r}
             (worker-util/profile (str "search: '" query-string "'")
               (js->clj (c.m/<? (.search infer-worker repo query-string nums-neighbors)) :keywordize-keys true))