Added support for local ONNX/model loading logic from worker.js

src/App.jsx

@@ -26,6 +26,7 @@ function App() {
   const [loadingMessage, setLoadingMessage] = useState("");
   const [progressItems, setProgressItems] = useState([]);
   const [isRunning, setIsRunning] = useState(false);
+  const [modelFiles, setModelFiles] = useState([]);
 
   // Inputs and outputs
   const [input, setInput] = useState("");
@@ -100,7 +101,6 @@ function App() {
     const onMessageReceived = (e) => {
       switch (e.data.status) {
         case "loading":
-          // Model file start load: add a new progress item to the list.
           setStatus("loading");
           setLoadingMessage(e.data.data);
           break;
@@ -110,7 +110,6 @@ function App() {
           break;
 
         case "progress":
-          // Model file progress: update one of the progress items.
           setProgressItems((prev) =>
             prev.map((item) => {
               if (item.file === e.data.file) {
@@ -122,53 +121,45 @@ function App() {
           break;
 
         case "done":
-          // Model file loaded: remove the progress item from the list.
           setProgressItems((prev) =>
             prev.filter((item) => item.file !== e.data.file),
           );
           break;
 
         case "ready":
-          // Pipeline ready: the worker is ready to accept messages.
           setStatus("ready");
           break;
 
         case "start":
-          {
-            // Start generation
-            setMessages((prev) => [
-              ...prev,
-              { role: "assistant", content: "" },
-            ]);
-          }
+          setMessages((prev) => [
+            ...prev,
+            { role: "assistant", content: "" },
+          ]);
           break;
 
         case "update":
-          {
-            // Generation update: update the output text.
-            // Parse messages
-            const { output, tps, numTokens } = e.data;
-            setTps(tps);
-            setNumTokens(numTokens);
-            setMessages((prev) => {
-              const cloned = [...prev];
-              const last = cloned.at(-1);
-              cloned[cloned.length - 1] = {
-                ...last,
-                content: last.content + output,
-              };
-              return cloned;
-            });
-          }
+          const { output, tps, numTokens } = e.data;
+          setTps(tps);
+          setNumTokens(numTokens);
+          setMessages((prev) => {
+            const cloned = [...prev];
+            const last = cloned.at(-1);
+            cloned[cloned.length - 1] = {
+              ...last,
+              content: last.content + output,
+            };
+            return cloned;
+          });
           break;
 
         case "complete":
-          // Generation complete: re-enable the "Generate" button
           setIsRunning(false);
           break;
 
         case "error":
-          setError(e.data.data);
+          setError(e.data.data || "Unknown error during model loading.");
+          setStatus(null);
+          setLoadingMessage("");
           break;
       }
     };
@@ -266,16 +257,36 @@ function App() {
               </div>
             )}
 
-            <button
-              className="border px-4 py-2 rounded-lg bg-blue-400 text-white hover:bg-blue-500 disabled:bg-blue-100 disabled:cursor-not-allowed select-none"
-              onClick={() => {
-                worker.current.postMessage({ type: "load" });
-                setStatus("loading");
-              }}
-              disabled={status !== null || error !== null}
-            >
-              Load model
-            </button>
+            <div className="flex flex-col items-center gap-2">
+              <label className="border px-4 py-2 rounded-lg bg-blue-100 text-blue-700 cursor-pointer hover:bg-blue-200">
+                Select model directory
+                <input
+                  type="file"
+                  webkitdirectory="true"
+                  directory="true"
+                  multiple
+                  className="hidden"
+                  onChange={e => {
+                    const files = Array.from(e.target.files);
+                    setModelFiles(files);
+                  }}
+                  disabled={status !== null || error !== null}
+                />
+              </label>
+              <button
+                className="border px-4 py-2 rounded-lg bg-blue-400 text-white hover:bg-blue-500 disabled:bg-blue-100 disabled:cursor-not-allowed select-none"
+                onClick={() => {
+                  worker.current.postMessage({ type: "load", files: modelFiles });
+                  setStatus("loading");
+                }}
+                disabled={status !== null || error !== null}
+              >
+                {modelFiles.length > 0 ? `Load selected model directory` : `Load default model`}
+              </button>
+              {modelFiles.length > 0 && (
+                <span className="text-xs text-gray-600">Selected files: {modelFiles.map(f => f.name).join(", ")}</span>
+              )}
+            </div>
           </div>
         </div>
       )}
@@ -376,7 +387,9 @@ function App() {
           }}
           onInput={(e) => setInput(e.target.value)}
         />
-        <label className="flex items-center px-2 py-2 bg-blue-100 text-blue-700 rounded cursor-pointer hover:bg-blue-200">
+        <label
+          className={`flex items-center px-2 py-2 bg-blue-100 text-blue-700 rounded cursor-pointer hover:bg-blue-200 ${status !== "ready" ? "opacity-50 cursor-not-allowed" : ""}`}
+        >
           📎 Attach
           <input
             type="file"

src/worker.js

@@ -30,18 +30,15 @@ async function check() {
  */
 class TextGenerationPipeline {
   static model_id = "onnx-community/granite-4.0-micro-ONNX-web";
-
   static async getInstance(progress_callback = null) {
     this.tokenizer ??= AutoTokenizer.from_pretrained(this.model_id, {
       progress_callback,
     });
-
     this.model ??= AutoModelForCausalLM.from_pretrained(this.model_id, {
       dtype: "q4f16",
       device: "webgpu",
       progress_callback,
     });
-
     return Promise.all([this.tokenizer, this.model]);
   }
 }
@@ -114,18 +111,16 @@ async function generate(messages) {
   });
 }
 
-async function load() {
+async function load(files = null) {
   self.postMessage({
     status: "loading",
-    data: "Loading model...",
+    data: files && files.length > 0 ? `Loading model from selected directory` : "Loading model...",
   });
 
   // Load the pipeline and save it for future use.
   const [tokenizer, model] = await TextGenerationPipeline.getInstance((x) => {
-    // We also add a progress callback to the pipeline so that we can
-    // track model loading.
     self.postMessage(x);
-  });
+  }, files);
 
   self.postMessage({
     status: "loading",
@@ -139,7 +134,7 @@ async function load() {
 }
 // Listen for messages from the main thread
 self.addEventListener("message", async (e) => {
-  const { type, data } = e.data;
+  const { type, data, files } = e.data;
 
   switch (type) {
     case "check":
@@ -147,7 +142,7 @@ self.addEventListener("message", async (e) => {
       break;
 
     case "load":
-      load();
+      await load(files || null);
       break;
 
     case "generate":