_source/src/hooks/useModel.js

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

export function useModel() {
  const workerRef = useRef(null);
  const audioUrlRef = useRef(null);
  const [status, setStatus] = useState("idle");
  const [message, setMessage] = useState("");
  const [progress, setProgress] = useState(null);
  const [audioUrl, setAudioUrl] = useState(null);
  const [audioInfo, setAudioInfo] = useState(null);
  const [error, setError] = useState(null);
  const [isLoaded, setIsLoaded] = useState(false);

  // Revoke a URL owned by this hook and forget it.
  const revokeCurrentAudioUrl = useCallback(() => {
    if (audioUrlRef.current) {
      URL.revokeObjectURL(audioUrlRef.current);
      audioUrlRef.current = null;
    }
  }, []);

  useEffect(() => {
    const worker = new Worker(new URL("../worker.js", import.meta.url), {
      type: "module",
    });

    worker.onmessage = (e) => {
      const { type, ...data } = e.data;
      switch (type) {
        case "status":
          setMessage(data.message);
          break;
        case "progress":
          setProgress(data);
          break;
        case "loaded":
          setIsLoaded(true);
          setStatus("ready");
          setProgress(null);
          break;
        case "audio": {
          // Revoke any previous URL owned by this hook before overwriting.
          if (audioUrlRef.current) URL.revokeObjectURL(audioUrlRef.current);
          const blob = new Blob([data.wavBuffer], { type: "audio/wav" });
          const url = URL.createObjectURL(blob);
          audioUrlRef.current = url;
          setAudioUrl(url);
          setAudioInfo({
            duration: data.duration,
            diffusionTime: data.diffusionTime,
            totalTime: data.totalTime,
            filename: `ace-step-${data.filenameStamp || Date.now()}.wav`,
          });
          setStatus("ready");
          setMessage("Generation complete!");
          break;
        }
        case "error":
          setError(data.message);
          setStatus("error");
          console.error("Worker error:", data.message, data.stack);
          break;
      }
    };

    workerRef.current = worker;
    return () => {
      worker.terminate();
      if (audioUrlRef.current) {
        URL.revokeObjectURL(audioUrlRef.current);
        audioUrlRef.current = null;
      }
    };
  }, []);

  const loadModel = useCallback(() => {
    setStatus("loading");
    setError(null);
    workerRef.current?.postMessage({ type: "load" });
  }, []);

  const generate = useCallback(({ caption, lyrics, duration, shift, numSteps }) => {
    setStatus("generating");
    setError(null);
    // Revoke the previous URL when user starts a new gen so the next "audio" message
    // doesn't compete with a still-displayed blob.
    revokeCurrentAudioUrl();
    setAudioUrl(null);
    setAudioInfo(null);
    workerRef.current?.postMessage({
      type: "generate",
      caption,
      lyrics,
      duration,
      shift,
      numSteps,
    });
  }, [revokeCurrentAudioUrl]);

  return {
    status,
    message,
    progress,
    audioUrl,
    audioInfo,
    error,
    isLoaded,
    loadModel,
    generate,
  };
}