Add error handling and message improvements

src/App.tsx

@@ -8,6 +8,7 @@ import { getModelsByPipeline } from './lib/huggingface'
 import TextGeneration from './components/pipelines/TextGeneration'
 import FeatureExtraction from './components/pipelines/FeatureExtraction'
 import ImageClassification from './components/pipelines/ImageClassification'
+import TextToSpeech from './components/pipelines/TextToSpeech'
 import Sidebar from './components/Sidebar'
 import ModelReadme from './components/ModelReadme'
 import { PipelineLayout } from './components/PipelineLayout'
@@ -18,13 +19,22 @@ function App() {
   const [isSidebarOpen, setIsSidebarOpen] = useState(false)
   const [isModalOpen, setIsModalOpen] = useState(false)
   const [isCodeModalOpen, setIsCodeModalOpen] = useState(false)
-  const { pipeline, setModels, setModelInfo, modelInfo, setIsFetching } =
-    useModel()
+  const {
+    pipeline,
+    setModels,
+    setModelInfo,
+    modelInfo,
+    setIsFetching,
+    setErrorText,
+    setStatus
+  } = useModel()
 
   useEffect(() => {
     setModelInfo(null)
     setModels([])
     setIsFetching(true)
+    setStatus('initiate')
+    setErrorText('')
 
     const fetchModels = async () => {
       try {
@@ -36,7 +46,14 @@ function App() {
       }
     }
     fetchModels()
-  }, [setModels, setModelInfo, setIsFetching, pipeline])
+  }, [
+    setModels,
+    setModelInfo,
+    setIsFetching,
+    setStatus,
+    setErrorText,
+    pipeline
+  ])
 
   return (
     <div className="relative min-h-screen flex flex-col bg-gradient-to-br from-blue-50 to-indigo-100">
@@ -67,6 +84,7 @@ function App() {
                 {pipeline === 'text-generation' && <TextGeneration />}
                 {pipeline === 'feature-extraction' && <FeatureExtraction />}
                 {pipeline === 'image-classification' && <ImageClassification />}
+                {pipeline === 'text-to-speech' && <TextToSpeech />}
               </div>
             </div>
           </main>

src/components/ModelInfo.tsx

@@ -25,7 +25,14 @@ const ModelInfo = () => {
     return num.toString()
   }
 
-  const { models, modelInfo, selectedQuantization, isFetching } = useModel()
+  const {
+    models,
+    status,
+    modelInfo,
+    selectedQuantization,
+    isFetching,
+    errorText
+  } = useModel()
 
   const ModelInfoSkeleton = () => (
     <div className="bg-gradient-to-r from-secondary to-accent px-3 py-3 rounded-lg border border-border space-y-3 animate-pulse w-4/5">
@@ -73,7 +80,7 @@ const ModelInfo = () => {
         {/* Compatibility Status */}
         {typeof modelInfo.isCompatible === 'boolean' && (
           <div className="shrink-0 ">
-            {modelInfo.isCompatible ? (
+            {modelInfo.isCompatible && status !== 'error' ? (
               <CheckCircle className="w-4 h-4 text-green-500" />
             ) : (
               <XCircle className="w-4 h-4 text-destructive" />
@@ -156,10 +163,11 @@ const ModelInfo = () => {
       <ModelLoader />
 
       {/* Incompatibility Message */}
-      {modelInfo.isCompatible === false && modelInfo.incompatibilityReason && (
+      {((modelInfo.isCompatible === false && modelInfo.incompatibilityReason) ||
+        errorText) && (
         <div className="bg-destructive/10 border border-destructive/20 rounded-md px-2 py-2">
           <p className="text-xs text-destructive whitespace-break-spaces">
-            {modelInfo.incompatibilityReason}
+            {errorText ? errorText : modelInfo.incompatibilityReason}
           </p>
         </div>
       )}

src/components/ModelLoader.tsx

@@ -2,12 +2,13 @@ import { useEffect, useCallback, useState } from 'react'
 import { ChevronDown, Loader2, X } from 'lucide-react'
 import { QuantizationType, WorkerMessage } from '../types'
 import { useModel } from '../contexts/ModelContext'
-import { getWorker } from '../lib/workerManager'
+import { getWorker, terminateWorker } from '../lib/workerManager'
 import { Alert, AlertDescription } from './ui/alert'
 
 const ModelLoader = () => {
   const [showAlert, setShowAlert] = useState(false)
   const [alertMessage, setAlertMessage] = useState<React.ReactNode>('')
+  const [lastModel, setLastModel] = useState<string | null>(null)
   const {
     modelInfo,
     selectedQuantization,
@@ -20,7 +21,8 @@ const ModelLoader = () => {
     setActiveWorker,
     pipeline,
     hasBeenLoaded,
-    setHasBeenLoaded
+    setHasBeenLoaded,
+    setErrorText
   } = useModel()
 
   useEffect(() => {
@@ -57,8 +59,10 @@ const ModelLoader = () => {
     }
 
     if (!hasBeenLoaded) {
+      setErrorText('')
       setStatus('initiate')
       setActiveWorker(newWorker)
+      setProgress(0)
     }
 
     const onMessageReceived = (e: MessageEvent<WorkerMessage>) => {
@@ -75,24 +79,25 @@ const ModelLoader = () => {
           output.file.startsWith('onnx')
         ) {
           setProgress(output.progress)
-          // setShowAlert(true)
-          // setAlertMessage(
-          //   <div className="flex items-center">
-          //     <Loader2 className="animate-spin h-4 w-4 mr-2" />
-          //     Loading Model
-          //   </div>
-          // )
         }
       } else if (status === 'error') {
         setStatus('error')
         const error = e.data.output
         console.error(error)
-        setAlertMessage(error.split('.')[0] + '. See console for details.')
+        const errText = error.split(' WASM error: ')[1]
+        setErrorText(errText)
         setShowAlert(true)
+        let time = 3000
+        if (!hasBeenLoaded)
+          setAlertMessage(error.split('.')[0] + '. See console for details.')
+        else {
+          setAlertMessage(`${errText}. Refresh the page and try again.`)
+          time = 5000
+        }
         setTimeout(() => {
           setShowAlert(false)
           setAlertMessage('')
-        }, 3000)
+        }, time)
       }
     }
 
@@ -100,7 +105,7 @@ const ModelLoader = () => {
 
     return () => {
       newWorker.removeEventListener('message', onMessageReceived)
-      // terminateWorker(pipeline);
+      terminateWorker(pipeline)
     }
   }, [
     pipeline,
@@ -110,7 +115,8 @@ const ModelLoader = () => {
     setStatus,
     setProgress,
     hasBeenLoaded,
-    setHasBeenLoaded
+    setHasBeenLoaded,
+    setErrorText
   ])
 
   useEffect(() => {
@@ -174,8 +180,10 @@ const ModelLoader = () => {
         {selectedQuantization && (
           <div className="flex justify-center">
             <button
-              className="w-32 py-2 px-4 bg-green-500 hover:bg-green-600 rounded-sm text-white font-medium disabled:opacity-50 disabled:cursor-not-allowed transition-colors text-sm inline-flex items-center text-center justify-center space-x-2"
-              disabled={hasBeenLoaded || status === 'loading'}
+              className={`w-32 py-2 px-4 ${status !== 'error' ? 'bg-green-500 hover:bg-green-600 cursor-pointer' : 'bg-red-500 hover:bg-red-600'} rounded-sm text-white font-medium disabled:opacity-50 disabled:cursor-not-allowed transition-colors text-sm inline-flex items-center text-center justify-center space-x-2`}
+              disabled={
+                hasBeenLoaded || status === 'loading' || status === 'error'
+              }
               onClick={loadModel}
             >
               {status === 'loading' && !hasBeenLoaded ? (
@@ -183,8 +191,10 @@ const ModelLoader = () => {
                   <Loader2 className="animate-spin h-4 w-4" />
                   <span>{progress.toFixed(0)}%</span>
                 </>
-              ) : (
+              ) : status !== 'error' ? (
                 <span>{!hasBeenLoaded ? 'Load Model' : 'Model Ready'}</span>
+              ) : (
+                <span>Error</span>
               )}
             </button>
           </div>

src/components/Sidebar.tsx

@@ -8,6 +8,7 @@ import FeatureExtractionConfig from './pipelines/FeatureExtractionConfig'
 import ZeroShotClassificationConfig from './pipelines/ZeroShotClassificationConfig'
 import ImageClassificationConfig from './pipelines/ImageClassificationConfig'
 import TextClassificationConfig from './pipelines/TextClassificationConfig'
+import TextToSpeechConfig from './pipelines/TextToSpeechConfig'
 import { Button } from '@/components/ui/button'
 import Tooltip from './Tooltip'
 
@@ -24,7 +25,7 @@ const Sidebar = ({
   setIsModalOpen,
   setIsCodeModalOpen
 }: SidebarProps) => {
-  const { pipeline, setPipeline } = useModel()
+  const { pipeline, setPipeline, modelInfo } = useModel()
 
   return (
     <>
@@ -95,13 +96,18 @@ const Sidebar = ({
               <ModelInfo />
               {/* Model README Button */}
               <div className="flex flex-row mt-2 space-x-4 ">
-                <Button variant="outline" onClick={() => setIsModalOpen(true)}>
+                <Button
+                  variant="outline"
+                  onClick={() => setIsModalOpen(true)}
+                  disabled={!modelInfo}
+                >
                   <FileText className="w-4 h-4 flex-shrink-0" />
                   <span>View README.md</span>
                 </Button>
                 <Button
                   variant="outline"
                   onClick={() => setIsCodeModalOpen(true)}
+                  disabled={!modelInfo}
                 >
                   <Code2 className="w-4 h-4 flex-shrink-0" />
                   <span>See Code</span>
@@ -119,6 +125,7 @@ const Sidebar = ({
               <ImageClassificationConfig />
             )}
             {pipeline === 'text-classification' && <TextClassificationConfig />}
+            {pipeline === 'text-to-speech' && <TextToSpeechConfig />}
           </div>
         </div>
       </div>

src/components/ui/button.tsx

@@ -1,37 +1,37 @@
-import * as React from "react"
-import { Slot } from "@radix-ui/react-slot"
-import { cva, type VariantProps } from "class-variance-authority"
+import * as React from 'react'
+import { Slot } from '@radix-ui/react-slot'
+import { cva, type VariantProps } from 'class-variance-authority'
 
-import { cn } from "@/lib/utils"
+import { cn } from '@/lib/utils'
 
 const buttonVariants = cva(
-  "inline-flex items-center justify-center gap-2 whitespace-nowrap rounded-md text-sm font-medium transition-all disabled:pointer-events-none disabled:opacity-50 [&_svg]:pointer-events-none [&_svg:not([class*='size-'])]:size-4 shrink-0 [&_svg]:shrink-0 outline-none focus-visible:border-ring focus-visible:ring-ring/50 focus-visible:ring-[3px] aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40 aria-invalid:border-destructive",
+  "cursor-pointer inline-flex items-center justify-center gap-2 whitespace-nowrap rounded-md text-sm font-medium transition-all disabled:pointer-events-none disabled:opacity-50 [&_svg]:pointer-events-none [&_svg:not([class*='size-'])]:size-4 shrink-0 [&_svg]:shrink-0 outline-none focus-visible:border-ring focus-visible:ring-ring/50 focus-visible:ring-[3px] aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40 aria-invalid:border-destructive",
   {
     variants: {
       variant: {
         default:
-          "bg-primary text-primary-foreground shadow-xs hover:bg-primary/90",
+          'bg-primary text-primary-foreground shadow-xs hover:bg-primary/90',
         destructive:
-          "bg-destructive text-white shadow-xs hover:bg-destructive/90 focus-visible:ring-destructive/20 dark:focus-visible:ring-destructive/40 dark:bg-destructive/60",
+          'bg-destructive text-white shadow-xs hover:bg-destructive/90 focus-visible:ring-destructive/20 dark:focus-visible:ring-destructive/40 dark:bg-destructive/60',
         outline:
-          "border bg-background shadow-xs hover:bg-accent hover:text-accent-foreground dark:bg-input/30 dark:border-input dark:hover:bg-input/50",
+          'border bg-background shadow-xs hover:bg-accent hover:text-accent-foreground dark:bg-input/30 dark:border-input dark:hover:bg-input/50',
         secondary:
-          "bg-secondary text-secondary-foreground shadow-xs hover:bg-secondary/80",
+          'bg-secondary text-secondary-foreground shadow-xs hover:bg-secondary/80',
         ghost:
-          "hover:bg-accent hover:text-accent-foreground dark:hover:bg-accent/50",
-        link: "text-primary underline-offset-4 hover:underline",
+          'hover:bg-accent hover:text-accent-foreground dark:hover:bg-accent/50',
+        link: 'text-primary underline-offset-4 hover:underline'
       },
       size: {
-        default: "h-9 px-4 py-2 has-[>svg]:px-3",
-        sm: "h-8 rounded-md gap-1.5 px-3 has-[>svg]:px-2.5",
-        lg: "h-10 rounded-md px-6 has-[>svg]:px-4",
-        icon: "size-9",
-      },
+        default: 'h-9 px-4 py-2 has-[>svg]:px-3',
+        sm: 'h-8 rounded-md gap-1.5 px-3 has-[>svg]:px-2.5',
+        lg: 'h-10 rounded-md px-6 has-[>svg]:px-4',
+        icon: 'size-9'
+      }
     },
     defaultVariants: {
-      variant: "default",
-      size: "default",
-    },
+      variant: 'default',
+      size: 'default'
+    }
   }
 )
 
@@ -41,11 +41,11 @@ function Button({
   size,
   asChild = false,
   ...props
-}: React.ComponentProps<"button"> &
+}: React.ComponentProps<'button'> &
   VariantProps<typeof buttonVariants> & {
     asChild?: boolean
   }) {
-  const Comp = asChild ? Slot : "button"
+  const Comp = asChild ? Slot : 'button'
 
   return (
     <Comp

src/contexts/ModelContext.tsx

@@ -25,6 +25,8 @@ interface ModelContextType {
   setIsFetching: (isFetching: boolean) => void
   hasBeenLoaded: boolean
   setHasBeenLoaded: (hasBeenLoaded: boolean) => void
+  errorText: string
+  setErrorText: (errorText: string) => void
 }
 
 const ModelContext = createContext<ModelContextType | undefined>(undefined)
@@ -42,6 +44,7 @@ export function ModelProvider({ children }: { children: React.ReactNode }) {
   const [activeWorker, setActiveWorker] = useState<Worker | null>(null)
   const [isFetching, setIsFetching] = useState(false)
   const [hasBeenLoaded, setHasBeenLoaded] = useState(false)
+  const [errorText, setErrorText] = useState('')
 
   // set progress to 0 when model is changed
   useEffect(() => {
@@ -68,7 +71,9 @@ export function ModelProvider({ children }: { children: React.ReactNode }) {
         isFetching,
         setIsFetching,
         hasBeenLoaded,
-        setHasBeenLoaded
+        setHasBeenLoaded,
+        errorText,
+        setErrorText
       }}
     >
       {children}