File size: 4,624 Bytes
ad5cef3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 |
import { useEffect, useCallback } from 'react'
import { ChevronDown, Loader } from 'lucide-react'
import { QuantizationType, WorkerMessage } from '../types'
import { useModel } from '../contexts/ModelContext'
import { getWorker } from '../lib/workerManager'
const ModelLoader = () => {
const {
modelInfo,
selectedQuantization,
setSelectedQuantization,
status,
progress,
setStatus,
setProgress,
activeWorker,
setActiveWorker,
pipeline
} = useModel()
useEffect(() => {
if (modelInfo.isCompatible && modelInfo.supportedQuantizations.length > 0) {
const quantizations = modelInfo.supportedQuantizations
let defaultQuant: QuantizationType = 'fp32'
if (quantizations.includes('int8')) {
defaultQuant = 'int8'
} else if (quantizations.includes('q8')) {
defaultQuant = 'q8'
} else if (quantizations.includes('q4')) {
defaultQuant = 'q4'
}
setSelectedQuantization(defaultQuant)
}
}, [
modelInfo.supportedQuantizations,
modelInfo.isCompatible,
setSelectedQuantization
])
useEffect(() => {
const newWorker = getWorker(pipeline)
if (!newWorker) {
return
}
setStatus('initiate')
setActiveWorker(newWorker)
const onMessageReceived = (e: MessageEvent<WorkerMessage>) => {
const { status, output } = e.data
if (status === 'ready') {
setStatus('ready')
} else if (status === 'loading' && output) {
setStatus('loading')
if (
output.progress &&
typeof output.file === 'string' &&
output.file.startsWith('onnx')
) {
setProgress(output.progress)
}
}
}
newWorker.addEventListener('message', onMessageReceived)
return () => {
newWorker.removeEventListener('message', onMessageReceived)
// terminateWorker(pipeline);
}
}, [pipeline, modelInfo.name, selectedQuantization, setActiveWorker, setStatus, setProgress])
const loadModel = useCallback(() => {
if (!modelInfo.name || !selectedQuantization) return
setStatus('loading')
const message = {
type: 'load',
model: modelInfo.name,
quantization: selectedQuantization
}
activeWorker?.postMessage(message)
}, [modelInfo.name, selectedQuantization, setStatus, activeWorker])
const ready: boolean = status === 'ready'
const busy: boolean = status === 'loading'
if (!modelInfo.isCompatible || modelInfo.supportedQuantizations.length === 0) {
return null
}
return (
<div className="space-y-3">
<hr className="border-gray-200" />
<div className="flex items-center justify-between space-x-4">
<div className="flex items-center space-x-2">
<span className="text-xs text-gray-600 font-medium">
Quantization:
</span>
<div className="relative">
<select
value={selectedQuantization || ''}
onChange={(e) =>
setSelectedQuantization(e.target.value as QuantizationType)
}
className="appearance-none bg-white border border-gray-300 rounded-md px-3 py-1 pr-8 text-xs text-gray-700 focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-blue-500"
>
<option value="">Select quantization</option>
{modelInfo.supportedQuantizations.map((quant) => (
<option key={quant} value={quant}>
{quant}
</option>
))}
</select>
<ChevronDown className="absolute right-2 top-1/2 transform -translate-y-1/2 w-3 h-3 text-gray-400 pointer-events-none" />
</div>
</div>
{selectedQuantization && (
<div className="flex justify-center">
<button
className="w-32 py-2 px-4 bg-green-500 hover:bg-green-600 rounded 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={(busy && !ready) || !selectedQuantization || ready}
onClick={loadModel}
>
{status === 'loading' && (
<>
<Loader className="animate-spin h-4 w-4" />
<span>{progress.toFixed(0)}%</span>
</>
)}
{!ready && !busy ? <span>Load Model</span> : !ready ? null : <span>Model Ready</span>}
</button>
</div>
)}
</div>
</div>
)
}
export default ModelLoader
|