whisper-word-timestamps/src/App.jsx

import { useEffect, useState, useRef, useCallback } from "react";

import Progress from "./components/Progress";
import MediaInput from "./components/MediaInput";
import Transcript from "./components/Transcript";
import LanguageSelector from "./components/LanguageSelector";

async function hasWebGPU() {
  if (!navigator.gpu) {
    return false;
  }
  try {
    const adapter = await navigator.gpu.requestAdapter();
    return !!adapter;
  } catch (e) {
    return false;
  }
}

function App() {
  // Create a reference to the worker object.
  const worker = useRef(null);

  // Model loading and progress
  const [status, setStatus] = useState(null);
  const [loadingMessage, setLoadingMessage] = useState("");
  const [progressItems, setProgressItems] = useState([]);

  const mediaInputRef = useRef(null);
  const [audio, setAudio] = useState(null);
  const [language, setLanguage] = useState("en");

  const [result, setResult] = useState(null);
  const [time, setTime] = useState(null);
  const [currentTime, setCurrentTime] = useState(0);

  const [device, setDevice] = useState("webgpu"); // Try use WebGPU first
  const [modelSize, setModelSize] = useState("gpu" in navigator ? 196 : 77); // WebGPU=196MB, WebAssembly=77MB
  useEffect(() => {
    hasWebGPU().then((result) => {
      setModelSize(result ? 196 : 77);
      setDevice(result ? "webgpu" : "wasm");
    });
  }, []);

  // We use the `useEffect` hook to setup the worker as soon as the `App` component is mounted.
  useEffect(() => {
    // Create the worker if it does not yet exist.
    worker.current ??= new Worker(new URL("./worker.js", import.meta.url), {
      type: "module",
    });

    // Create a callback function for messages from the worker thread.
    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;

        case "initiate":
          setProgressItems((prev) => [...prev, e.data]);
          break;

        case "progress":
          // Model file progress: update one of the progress items.
          setProgressItems((prev) =>
            prev.map((item) => {
              if (item.file === e.data.file) {
                return { ...item, ...e.data };
              }
              return item;
            }),
          );
          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 "complete":
          setResult(e.data.result);
          setTime(e.data.time);
          setStatus("ready");
          break;
      }
    };

    // Attach the callback function as an event listener.
    worker.current.addEventListener("message", onMessageReceived);

    // Define a cleanup function for when the component is unmounted.
    return () => {
      worker.current.removeEventListener("message", onMessageReceived);
    };
  }, []);

  const handleClick = useCallback(() => {
    setResult(null);
    setTime(null);
    if (status === null) {
      setStatus("loading");
      worker.current.postMessage({ type: "load", data: { device } });
    } else {
      setStatus("running");
      worker.current.postMessage({
        type: "run",
        data: { audio, language },
      });
    }
  }, [status, audio, language, device]);

  return (
    <div className="w-screen h-screen text-gray-800 dark:text-gray-200 bg-white dark:bg-gray-900 ">

      <div className="flex flex-col mx-auto items justify-end max-w-[560px] h-full">

        {status === "loading" && (

          <div className="flex justify-center items-center fixed w-screen h-screen bg-black z-10 bg-opacity-[92%] top-0 left-0">

            <div className="w-[500px]">

              <p className="text-center mb-1 text-white text-md">

                {loadingMessage}

              </p>

              {progressItems.map(({ file, progress, total }, i) => (

                <Progress

                  key={i}

                  text={file}

                  percentage={progress}

                  total={total}

                />

              ))}

            </div>

          </div>

        )}

        <div className="h-full flex justify-center items-center flex-col relative">

          <div className="flex flex-col items-center mb-1 text-center">

            <h1 className="text-5xl font-bold mb-2">Whisper Timestamped</h1>

            <h2 className="text-xl font-semibold">

              In-browser speech recognition w/ word-level timestamps

            </h2>

          </div>



          <div className="w-full min-h-[220px] flex flex-col justify-center items-center p-2">

            {!audio && (

              <p className="mb-2">

                You are about to download{" "}

                <a

                  href="https://huggingface.co/onnx-community/whisper-base_timestamped"

                  target="_blank"

                  rel="noreferrer"

                  className="font-medium underline"

                >

                  whisper-base (timestamped)

                </a>

                , a 73 million parameter speech recognition model with the

                ability to generate word-level timestamps across 100 different

                languages. Once loaded, the model ({modelSize}&nbsp;MB) will be

                cached and reused when you revisit the page.

                <br />

                <br />

                Everything runs locally in your browser using{" "}

                <a

                  href="https://huggingface.co/docs/transformers.js"

                  target="_blank"

                  rel="noreferrer"

                  className="underline"

                >

                  🤗&nbsp;Transformers.js

                </a>{" "}

                and ONNX Runtime Web, meaning no API calls are made to a server

                for inference. You can even disconnect from the internet after

                the model has loaded!

              </p>

            )}



            <div className="flex flex-col w-full m-3">

              <span className="text-sm mb-0.5">Input audio/video</span>

              <MediaInput

                ref={mediaInputRef}

                className="flex items-center border rounded-md cursor-pointer min-h-[100px] max-h-[500px] overflow-hidden"

                onInputChange={(result) => setAudio(result)}

                onTimeUpdate={(time) => setCurrentTime(time)}

              />

            </div>



            <div className="relative w-full flex justify-center items-center">

              <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 cursor-pointer"

                onClick={handleClick}

                disabled={

                  status === "running" || (status !== null && audio === null)

                }

              >

                {status === null

                  ? "Load model"

                  : status === "running"

                    ? "Running..."

                    : "Run model"}

              </button>



              {status !== null && (

                <div className="absolute right-0 bottom-0">

                  <span className="text-xs">Language:</span>

                  <br />

                  <LanguageSelector

                    className="border rounded-lg p-1 max-w-[100px] dark:bg-gray-800"

                    language={language}

                    setLanguage={setLanguage}

                  />

                </div>

              )}

            </div>



            {result && time && (

              <>

                <div className="w-full mt-4 border rounded-md">

                  <Transcript

                    className="p-2 max-h-[200px] overflow-y-auto scrollbar-thin select-none"

                    transcript={result}

                    currentTime={currentTime}

                    setCurrentTime={(time) => {

                      setCurrentTime(time);

                      mediaInputRef.current.setMediaTime(time);

                    }}

                  />

                </div>

                <p className="text-sm text-gray-600 dark:text-gray-300 text-end p-1">

                  Generation time:{" "}

                  <span className="text-gray-800 dark:text-gray-200 font-semibold">

                    {time.toFixed(2)}ms

                  </span>

                </p>

              </>

            )}

          </div>

        </div>

      </div>
    </div>
  );
}

export default App;