import { useState, useEffect } from 'react'; import { Download, Trash2, X, Check, ChevronDown, ChevronUp } from 'lucide-react'; import { Button } from '../../ui/button'; import { useConfig } from '../../ConfigContext'; import { listModels, downloadModel, getDownloadProgress, cancelDownload as cancelDownloadApi, deleteModel as deleteModelApi, type WhisperModelResponse, type DownloadProgress, } from '../../../api'; import { defineMessages, useIntl } from '../../../i18n'; const i18n = defineMessages({ gpuAcceleration: { id: 'localModelManager.gpuAcceleration', defaultMessage: 'Supports GPU acceleration (CUDA for NVIDIA, Metal for Apple Silicon). GPU features must be enabled at build time for hardware acceleration.', }, recommended: { id: 'localModelManager.recommended', defaultMessage: 'Recommended', }, active: { id: 'localModelManager.active', defaultMessage: 'Active', }, recommendedForHardware: { id: 'localModelManager.recommendedForHardware', defaultMessage: 'Recommended for your hardware', }, downloaded: { id: 'localModelManager.downloaded', defaultMessage: 'Downloaded', }, download: { id: 'localModelManager.download', defaultMessage: 'Download', }, deleteConfirm: { id: 'localModelManager.deleteConfirm', defaultMessage: 'Delete this model? You can re-download it later.', }, showRecommendedOnly: { id: 'localModelManager.showRecommendedOnly', defaultMessage: 'Show recommended only', }, showAllModels: { id: 'localModelManager.showAllModels', defaultMessage: 'Show all models ({count} more)', }, noModels: { id: 'localModelManager.noModels', defaultMessage: 'No models available', }, }); const LOCAL_WHISPER_MODEL_CONFIG_KEY = 'LOCAL_WHISPER_MODEL'; const formatBytes = (bytes: number): string => { if (bytes < 1024) return `${bytes}B`; if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(0)}KB`; if (bytes < 1024 * 1024 * 1024) return `${(bytes / (1024 * 1024)).toFixed(0)}MB`; return `${(bytes / (1024 * 1024 * 1024)).toFixed(1)}GB`; }; const capitalize = (str: string): string => { return str.charAt(0).toUpperCase() + str.slice(1); }; export const LocalModelManager = () => { const intl = useIntl(); const [models, setModels] = useState([]); const [downloads, setDownloads] = useState>(new Map()); const [selectedModelId, setSelectedModelId] = useState(null); const [showAllModels, setShowAllModels] = useState(false); const { read, upsert } = useConfig(); useEffect(() => { loadModels(); loadSelectedModel(); // eslint-disable-next-line react-hooks/exhaustive-deps }, []); const loadSelectedModel = async () => { try { const value = await read(LOCAL_WHISPER_MODEL_CONFIG_KEY, false); if (value && typeof value === 'string') { setSelectedModelId(value); } else { setSelectedModelId(null); } } catch (error) { console.error('Failed to load selected model:', error); setSelectedModelId(null); } }; const selectModel = async (modelId: string) => { await upsert(LOCAL_WHISPER_MODEL_CONFIG_KEY, modelId, false); setSelectedModelId(modelId); }; const loadModels = async () => { try { const response = await listModels(); if (response.data) { setModels(response.data); } } catch (error) { console.error('Failed to load models:', error); } }; const startDownload = async (modelId: string) => { try { await downloadModel({ path: { model_id: modelId } }); pollDownloadProgress(modelId); } catch (error) { console.error('Failed to start download:', error); } }; const pollDownloadProgress = (modelId: string) => { const interval = setInterval(async () => { try { const response = await getDownloadProgress({ path: { model_id: modelId } }); if (response.data) { const progress = response.data; setDownloads((prev) => new Map(prev).set(modelId, progress)); if (progress.status === 'completed') { clearInterval(interval); await loadModels(); // Refresh model list // Backend auto-selects, but also update frontend state await loadSelectedModel(); } else if (progress.status === 'failed') { clearInterval(interval); await loadModels(); } } else { clearInterval(interval); } } catch { clearInterval(interval); } }, 500); }; const cancelDownload = async (modelId: string) => { try { await cancelDownloadApi({ path: { model_id: modelId } }); setDownloads((prev) => { const next = new Map(prev); next.delete(modelId); return next; }); loadModels(); } catch (error) { console.error('Failed to cancel download:', error); } }; const deleteModel = async (modelId: string) => { if (!window.confirm(intl.formatMessage(i18n.deleteConfirm))) return; try { await deleteModelApi({ path: { model_id: modelId } }); if (selectedModelId === modelId) { await upsert(LOCAL_WHISPER_MODEL_CONFIG_KEY, '', false); setSelectedModelId(null); } loadModels(); } catch (error) { console.error('Failed to delete model:', error); } }; const hasDownloadedNonRecommended = models.some( (model) => model.downloaded && !model.recommended ); const displayedModels = showAllModels || hasDownloadedNonRecommended ? models : models.filter((m) => m.recommended); const hasNonRecommendedModels = models.some((m) => !m.recommended); const showToggleButton = hasNonRecommendedModels && !hasDownloadedNonRecommended; return (

{intl.formatMessage(i18n.gpuAcceleration)}

{displayedModels.map((model) => { const progress = downloads.get(model.id); const isDownloading = progress?.status === 'downloading'; const isSelected = selectedModelId === model.id; const canSelect = model.downloaded && !isDownloading; return (

{canSelect && ( selectModel(model.id)} className="cursor-pointer" /> )}

{capitalize(model.id)}

{model.size_mb}MB {model.recommended && ( {intl.formatMessage(i18n.recommended)} )} {isSelected && ( {intl.formatMessage(i18n.active)} )}

{model.description}

{model.recommended && (

{intl.formatMessage(i18n.recommendedForHardware)}

)}

{model.downloaded ? ( <>

{intl.formatMessage(i18n.downloaded)}

) : isDownloading ? ( <>

{progress.progress_percent.toFixed(0)}%

) : ( )}

{isDownloading && progress && (

{formatBytes(progress.bytes_downloaded)} / {formatBytes(progress.total_bytes)} {progress.speed_bps && {formatBytes(progress.speed_bps)}/s}

)} {progress?.status === 'failed' && progress.error && (

{progress.error}

)}

); })}

{showToggleButton && ( )} {models.length === 0 && (

{intl.formatMessage(i18n.noModels)}

)}

); };