import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'; import { ChevronDown, ChevronRight, ChevronUp, CircleCheck, CircleX, Download, ExternalLink, FolderOpen, HardDrive, Heart, Loader2, RotateCcw, Scale, Trash2, Unplug, X, } from 'lucide-react'; import { useCallback, useMemo, useState } from 'react'; import { Trans, useTranslation } from 'react-i18next'; import { AlertDialog, AlertDialogAction, AlertDialogCancel, AlertDialogContent, AlertDialogDescription, AlertDialogFooter, AlertDialogHeader, AlertDialogTitle, } from '@/components/ui/alert-dialog'; import { Badge } from '@/components/ui/badge'; import { Button } from '@/components/ui/button'; import { Dialog, DialogContent, DialogDescription, DialogHeader, DialogTitle, } from '@/components/ui/dialog'; import { Progress } from '@/components/ui/progress'; import { useToast } from '@/components/ui/use-toast'; import { apiClient } from '@/lib/api/client'; import type { ActiveDownloadTask, HuggingFaceModelInfo, ModelStatus } from '@/lib/api/types'; import { useModelDownloadToast } from '@/lib/hooks/useModelDownloadToast'; import { usePlatform } from '@/platform/PlatformContext'; import { useServerStore } from '@/stores/serverStore'; async function fetchHuggingFaceModelInfo(repoId: string): Promise { const response = await fetch(`https://huggingface.co/api/models/${repoId}`); if (!response.ok) throw new Error(`Failed to fetch model info: ${response.status}`); return response.json(); } const MODEL_DESCRIPTIONS: Record = { 'qwen-tts-1.7B': 'High-quality multilingual TTS by Alibaba. Supports 10 languages with natural prosody and voice cloning from short reference audio.', 'qwen-tts-0.6B': 'Lightweight version of Qwen TTS. Same language support with faster inference, ideal for lower-end hardware.', luxtts: 'Lightweight ZipVoice-based TTS designed for high quality voice cloning and 48kHz speech generation at speeds exceeding 150x realtime.', 'chatterbox-tts': 'Production-grade open source TTS by Resemble AI. Supports 23 languages with voice cloning and emotion exaggeration control.', 'chatterbox-turbo': 'Streamlined 350M parameter TTS by Resemble AI. High-quality English speech with less compute and VRAM than larger models.', 'tada-1b': 'HumeAI TADA 1B — English speech-language model built on Llama 3.2 1B. Generates 700s+ of coherent audio with synchronized text-acoustic alignment.', 'tada-3b-ml': 'HumeAI TADA 3B Multilingual — built on Llama 3.2 3B. Supports 10 languages with high-fidelity voice cloning via text-acoustic dual alignment.', kokoro: 'Kokoro 82M by hexgrad. Tiny 82M-parameter TTS that runs at CPU realtime. Supports 8 languages with pre-built voice styles. Apache 2.0 licensed.', 'qwen-custom-voice-1.7B': 'Qwen3-TTS CustomVoice 1.7B by Alibaba. 9 premium preset voices with instruct-based style control for tone, emotion, and prosody. Supports 10 languages.', 'qwen-custom-voice-0.6B': 'Qwen3-TTS CustomVoice 0.6B by Alibaba. Lightweight version with the same 9 preset voices and instruct control. Faster inference for lower-end hardware.', 'whisper-base': 'Smallest Whisper model (74M parameters). Fast transcription with moderate accuracy.', 'whisper-small': 'Whisper Small (244M parameters). Good balance of speed and accuracy for transcription.', 'whisper-medium': 'Whisper Medium (769M parameters). Higher accuracy transcription at moderate speed.', 'whisper-large': 'Whisper Large (1.5B parameters). Best accuracy for speech-to-text across multiple languages.', 'whisper-turbo': 'Whisper Large v3 Turbo. Pruned for significantly faster inference while maintaining near-large accuracy.', 'qwen3-0.6b': 'Qwen3 0.6B — smallest of the Qwen3 instruct family. Very fast on CPU, runs at ~400 MB quantized on Apple Silicon. Good for dictation refinement and short completions.', 'qwen3-1.7b': 'Qwen3 1.7B — balanced size and quality. Handles subtle self-corrections and technical vocabulary better than the 0.6B. Runs at ~1.1 GB quantized on Apple Silicon.', 'qwen3-4b': 'Qwen3 4B — highest quality local refinement and longer-form reasoning. ~2.5 GB quantized on Apple Silicon, ~8 GB at full precision on PyTorch.', }; function formatDownloads(n: number): string { if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`; if (n >= 1_000) return `${(n / 1_000).toFixed(1)}k`; return n.toString(); } function formatLicense(license: string): string { const map: Record = { 'apache-2.0': 'Apache 2.0', mit: 'MIT', 'cc-by-4.0': 'CC BY 4.0', 'cc-by-sa-4.0': 'CC BY-SA 4.0', 'cc-by-nc-4.0': 'CC BY-NC 4.0', 'openrail++': 'OpenRAIL++', openrail: 'OpenRAIL', }; return map[license] || license; } function formatPipelineTag(tag: string): string { return tag .split('-') .map((w) => w.charAt(0).toUpperCase() + w.slice(1)) .join(' '); } function formatBytes(bytes: number): string { if (bytes === 0) return '0 B'; const k = 1024; const sizes = ['B', 'KB', 'MB', 'GB']; const i = Math.floor(Math.log(bytes) / Math.log(k)); return `${(bytes / k ** i).toFixed(1)} ${sizes[i]}`; } export function ModelManagement() { const { t } = useTranslation(); const { toast } = useToast(); const queryClient = useQueryClient(); const platform = usePlatform(); const customModelsDir = useServerStore((state) => state.customModelsDir); const setCustomModelsDir = useServerStore((state) => state.setCustomModelsDir); const [migrating, setMigrating] = useState(false); const [migrationProgress, setMigrationProgress] = useState<{ current: number; total: number; progress: number; filename?: string; status: string; } | null>(null); const [pendingMigrateDir, setPendingMigrateDir] = useState(null); const [downloadingModel, setDownloadingModel] = useState(null); const [downloadingDisplayName, setDownloadingDisplayName] = useState(null); const [consoleOpen, setConsoleOpen] = useState(false); const [dismissedErrors, setDismissedErrors] = useState>(new Set()); const [localErrors, setLocalErrors] = useState>(new Map()); // Modal state const [selectedModel, setSelectedModel] = useState(null); const [detailOpen, setDetailOpen] = useState(false); const { data: modelStatus, isLoading } = useQuery({ queryKey: ['modelStatus'], queryFn: async () => { const result = await apiClient.getModelStatus(); return result; }, refetchInterval: 5000, }); const { data: cacheDir } = useQuery({ queryKey: ['modelsCacheDir'], queryFn: () => apiClient.getModelsCacheDir(), staleTime: 1000 * 60 * 5, }); const { data: activeTasks } = useQuery({ queryKey: ['activeTasks'], queryFn: () => apiClient.getActiveTasks(), refetchInterval: (query) => { const data = query.state.data; const hasActive = data?.downloads.some((d) => d.status === 'downloading'); return hasActive ? 1000 : 5000; }, }); // HuggingFace model card query - only fetches when modal is open and model has a repo ID const { data: hfModelInfo, isLoading: hfLoading } = useQuery({ queryKey: ['hfModelInfo', selectedModel?.hf_repo_id], queryFn: () => fetchHuggingFaceModelInfo(selectedModel!.hf_repo_id!), enabled: detailOpen && !!selectedModel?.hf_repo_id, staleTime: 1000 * 60 * 30, // Cache for 30 minutes retry: 1, }); // Build a map of errored downloads for quick lookup, excluding dismissed ones const erroredDownloads = new Map(); if (activeTasks?.downloads) { for (const dl of activeTasks.downloads) { if (dl.status === 'error' && !dismissedErrors.has(dl.model_name)) { const localErr = localErrors.get(dl.model_name); erroredDownloads.set(dl.model_name, localErr ? { ...dl, error: localErr } : dl); } } } for (const [modelName, error] of localErrors) { if (!erroredDownloads.has(modelName) && !dismissedErrors.has(modelName)) { erroredDownloads.set(modelName, { model_name: modelName, status: 'error', started_at: new Date().toISOString(), error, }); } } const errorCount = erroredDownloads.size; // Build progress map from active tasks for inline display const downloadProgressMap = useMemo(() => { const map = new Map(); if (activeTasks?.downloads) { for (const dl of activeTasks.downloads) { if (dl.status === 'downloading') { map.set(dl.model_name, dl); } } } return map; }, [activeTasks]); const handleDownloadComplete = useCallback(() => { setDownloadingModel(null); setDownloadingDisplayName(null); queryClient.invalidateQueries({ queryKey: ['modelStatus'] }); queryClient.invalidateQueries({ queryKey: ['activeTasks'] }); }, [queryClient]); const handleDownloadError = useCallback( (error: string) => { if (downloadingModel) { setLocalErrors((prev) => new Map(prev).set(downloadingModel, error)); setConsoleOpen(true); } setDownloadingModel(null); setDownloadingDisplayName(null); queryClient.invalidateQueries({ queryKey: ['activeTasks'] }); }, [queryClient, downloadingModel], ); useModelDownloadToast({ modelName: downloadingModel || '', displayName: downloadingDisplayName || '', enabled: !!downloadingModel && !!downloadingDisplayName, onComplete: handleDownloadComplete, onError: handleDownloadError, }); const [deleteDialogOpen, setDeleteDialogOpen] = useState(false); const [modelToDelete, setModelToDelete] = useState<{ name: string; displayName: string; sizeMb?: number; } | null>(null); const handleDownload = async (modelName: string) => { setDismissedErrors((prev) => { const next = new Set(prev); next.delete(modelName); return next; }); const model = modelStatus?.models.find((m) => m.model_name === modelName); const displayName = model?.display_name || modelName; try { await apiClient.triggerModelDownload(modelName); setDownloadingModel(modelName); setDownloadingDisplayName(displayName); queryClient.invalidateQueries({ queryKey: ['modelStatus'] }); queryClient.invalidateQueries({ queryKey: ['activeTasks'] }); } catch (error) { setDownloadingModel(null); setDownloadingDisplayName(null); toast({ title: t('models.toast.downloadFailed'), description: error instanceof Error ? error.message : t('common.unknownError'), variant: 'destructive', }); } }; const cancelMutation = useMutation({ mutationFn: (modelName: string) => apiClient.cancelDownload(modelName), onSuccess: async () => { await queryClient.invalidateQueries({ queryKey: ['modelStatus'], refetchType: 'all' }); await queryClient.invalidateQueries({ queryKey: ['activeTasks'], refetchType: 'all' }); }, }); const handleCancel = (modelName: string) => { const prevDismissed = dismissedErrors; const prevLocalErrors = localErrors; const prevDownloadingModel = downloadingModel; const prevDownloadingDisplayName = downloadingDisplayName; setDismissedErrors((prev) => new Set(prev).add(modelName)); setLocalErrors((prev) => { const next = new Map(prev); next.delete(modelName); return next; }); if (downloadingModel === modelName) { setDownloadingModel(null); setDownloadingDisplayName(null); } cancelMutation.mutate(modelName, { onError: () => { setDismissedErrors(prevDismissed); setLocalErrors(prevLocalErrors); setDownloadingModel(prevDownloadingModel); setDownloadingDisplayName(prevDownloadingDisplayName); toast({ title: t('models.toast.cancelFailed'), description: t('models.toast.cancelFailedDescription'), variant: 'destructive', }); }, }); }; const clearAllMutation = useMutation({ mutationFn: () => apiClient.clearAllTasks(), onSuccess: async () => { setDismissedErrors(new Set()); setLocalErrors(new Map()); setDownloadingModel(null); setDownloadingDisplayName(null); await queryClient.invalidateQueries({ queryKey: ['modelStatus'], refetchType: 'all' }); await queryClient.invalidateQueries({ queryKey: ['activeTasks'], refetchType: 'all' }); }, }); const deleteMutation = useMutation({ mutationFn: async (modelName: string) => { const result = await apiClient.deleteModel(modelName); return result; }, onSuccess: async () => { toast({ title: t('models.toast.deleted'), description: t('models.toast.deletedDescription', { name: modelToDelete?.displayName || t('models.defaultName'), }), }); setDeleteDialogOpen(false); setModelToDelete(null); setDetailOpen(false); setSelectedModel(null); await queryClient.invalidateQueries({ queryKey: ['modelStatus'], refetchType: 'all' }); await queryClient.refetchQueries({ queryKey: ['modelStatus'] }); }, onError: (error: Error) => { toast({ title: t('models.toast.deleteFailed'), description: error.message, variant: 'destructive', }); }, }); const unloadMutation = useMutation({ mutationFn: async (modelName: string) => { return await apiClient.unloadModel(modelName); }, onSuccess: async (_data, modelName) => { toast({ title: t('models.toast.unloaded'), description: t('models.toast.unloadedDescription', { name: modelName }), }); await queryClient.invalidateQueries({ queryKey: ['modelStatus'], refetchType: 'all' }); await queryClient.refetchQueries({ queryKey: ['modelStatus'] }); }, onError: (error: Error) => { toast({ title: t('models.toast.unloadFailed'), description: error.message, variant: 'destructive', }); }, }); const formatSize = (sizeMb?: number): string => { if (!sizeMb) return t('models.unknownSize'); if (sizeMb < 1024) return `${sizeMb.toFixed(1)} MB`; return `${(sizeMb / 1024).toFixed(2)} GB`; }; const getModelState = (model: ModelStatus) => { const isDownloading = (model.downloading || downloadingModel === model.model_name) && !erroredDownloads.has(model.model_name) && !dismissedErrors.has(model.model_name); const hasError = erroredDownloads.has(model.model_name); return { isDownloading, hasError }; }; const openModelDetail = (model: ModelStatus) => { setSelectedModel(model); setDetailOpen(true); }; const voiceModels = modelStatus?.models.filter( (m) => m.model_name.startsWith('qwen-tts') || m.model_name.startsWith('qwen-custom-voice') || m.model_name.startsWith('luxtts') || m.model_name.startsWith('chatterbox') || m.model_name.startsWith('tada') || m.model_name.startsWith('kokoro'), ) ?? []; const whisperModels = modelStatus?.models.filter((m) => m.model_name.startsWith('whisper')) ?? []; const llmModels = modelStatus?.models.filter((m) => m.model_name.startsWith('qwen3-')) ?? []; // Build sections const sections: { label: string; models: ModelStatus[] }[] = [ { label: t('models.sections.voiceGeneration'), models: voiceModels }, { label: t('models.sections.transcription'), models: whisperModels }, { label: t('models.sections.languageModels'), models: llmModels }, ]; // Get detail modal state for selected model const selectedState = selectedModel ? getModelState(selectedModel) : null; const selectedError = selectedModel ? erroredDownloads.get(selectedModel.model_name) : undefined; // Keep selectedModel data fresh from query results const freshSelectedModel = selectedModel && modelStatus ? modelStatus.models.find((m) => m.model_name === selectedModel.model_name) || selectedModel : selectedModel; // Derive license from HF data const license = hfModelInfo?.cardData?.license || hfModelInfo?.tags?.find((tag) => tag.startsWith('license:'))?.replace('license:', ''); return (
{/* Header */}

{t('models.title')}

{t('models.subtitle')}

{/* Model storage location */} {platform.metadata.isTauri && cacheDir && (
{t('models.storage.location')}

{cacheDir.path}

{customModelsDir && ( )}
)} {/* Model list */} {isLoading ? (
) : modelStatus ? (
{sections.map((section) => (

{section.label}

{section.models.map((model) => { const { isDownloading, hasError } = getModelState(model); return ( ); })}
))} {/* Error console */} {errorCount > 0 && (
{consoleOpen && (
{Array.from(erroredDownloads.entries()).map(([modelName, dl]) => (
[error]{' '} {modelName} {dl.error ? ( <> {': '} {dl.error} ) : ( <> {': '} {t('models.problems.noDetails')} )}
{t('models.problems.startedAt', { time: new Date(dl.started_at).toLocaleString(), })}
))}
)}
)}
) : null} {/* Model Detail Modal */} {freshSelectedModel && ( <> {freshSelectedModel.display_name} {freshSelectedModel.hf_repo_id ? ( {freshSelectedModel.hf_repo_id} ) : ( freshSelectedModel.model_name )}
{/* Status badges */}
{freshSelectedModel.loaded && ( {t('models.status.loaded')} )} {selectedState?.hasError && ( {t('common.error')} )}
{/* HuggingFace model card info */} {hfLoading && freshSelectedModel.hf_repo_id && (
{t('models.detail.loadingInfo')}
)} {/* Description */} {MODEL_DESCRIPTIONS[freshSelectedModel.model_name] && (

{MODEL_DESCRIPTIONS[freshSelectedModel.model_name]}

)} {hfModelInfo && (
{/* Pipeline tag + author */}
{hfModelInfo.pipeline_tag && ( {formatPipelineTag(hfModelInfo.pipeline_tag)} )} {hfModelInfo.library_name && ( {hfModelInfo.library_name} )} {hfModelInfo.author && ( {t('models.detail.byAuthor', { author: hfModelInfo.author })} )}
{/* Stats row */}
{formatDownloads(hfModelInfo.downloads)} {formatDownloads(hfModelInfo.likes)} {license && ( {formatLicense(license)} )}
{/* Languages */} {hfModelInfo.cardData?.language && hfModelInfo.cardData.language.length > 0 && (
{hfModelInfo.cardData.language.length > 10 ? t('models.detail.languagesCount', { count: hfModelInfo.cardData.language.length, }) : t('models.detail.languagesList', { list: hfModelInfo.cardData.language.join(', '), })}
)}
)} {/* Disk size */} {freshSelectedModel.downloaded && freshSelectedModel.size_mb && (
{t('models.detail.onDisk', { size: formatSize(freshSelectedModel.size_mb) })}
)} {/* Error detail */} {selectedError?.error && (
{selectedError.error}
)} {/* Actions */}
{selectedState?.hasError ? ( <> ) : selectedState?.isDownloading ? ( <>
{(() => { const dl = freshSelectedModel ? downloadProgressMap.get(freshSelectedModel.model_name) : undefined; const pct = dl?.progress ?? 0; const hasProgress = dl && dl.total && dl.total > 0; return ( <>
{hasProgress ? `${formatBytes(dl.current ?? 0)} / ${formatBytes(dl.total!)} (${pct.toFixed(1)}%)` : dl?.filename || t('models.progress.connectingHf')}
); })()}
) : freshSelectedModel.downloaded ? (
{freshSelectedModel.loaded && ( )}
) : ( )}
)}
{/* Delete Confirmation Dialog */} {t('models.deleteDialog.title')} }} /> {modelToDelete?.sizeMb && ( <> {' '} {t('models.deleteDialog.sizeNote', { size: formatSize(modelToDelete.sizeMb) })} )} {t('common.cancel')} { if (modelToDelete) { deleteMutation.mutate(modelToDelete.name); } }} disabled={deleteMutation.isPending} className="bg-destructive text-destructive-foreground hover:bg-destructive/90" > {deleteMutation.isPending ? ( <> {t('models.deleteDialog.deleting')} ) : ( t('common.delete') )} {/* Migration confirmation dialog */} !open && setPendingMigrateDir(null)} > {t('models.migrateDialog.title')} {t('models.migrateDialog.description')}
{pendingMigrateDir}
{t('common.cancel')} { if (!pendingMigrateDir) return; const newDir = pendingMigrateDir; setPendingMigrateDir(null); setMigrating(true); setMigrationProgress({ current: 0, total: 0, progress: 0, status: 'downloading', filename: t('models.migrateDialog.preparing'), }); try { // Start the migration (background task) const migrationResult = await apiClient.migrateModels(newDir); // If no models to migrate, warn user and skip the change if (migrationResult.moved === 0) { setMigrating(false); setMigrationProgress(null); toast({ title: t('models.toast.noModelsToMigrate'), description: t('models.toast.noModelsToMigrateDescription'), }); setPendingMigrateDir(null); return; } // Connect to SSE for progress await new Promise((resolve, reject) => { const es = new EventSource(apiClient.getMigrationProgressUrl()); es.onmessage = (event) => { try { const data = JSON.parse(event.data); setMigrationProgress(data); if (data.status === 'complete') { es.close(); resolve(); } else if (data.status === 'error') { es.close(); reject(new Error(data.error || t('models.toast.migrationFailed'))); } } catch { /* ignore parse errors */ } }; es.onerror = () => { es.close(); reject(new Error(t('models.toast.migrationConnectionLost'))); }; }); setCustomModelsDir(newDir); setMigrationProgress({ current: 1, total: 1, progress: 100, status: 'complete', filename: t('models.migrateDialog.restartingServer'), }); await platform.lifecycle.restartServer(newDir); queryClient.invalidateQueries(); toast({ title: t('models.toast.migrated') }); } catch (e) { toast({ title: t('models.toast.migrationFailed'), description: e instanceof Error ? e.message : t('models.toast.migrationFailedGeneric'), variant: 'destructive', }); } finally { setMigrating(false); setMigrationProgress(null); } }} > {t('models.migrateDialog.action')}
{/* Migration progress overlay */} {migrating && migrationProgress && (

{t('models.migrate.title')}

{migrationProgress.status === 'complete' ? t('models.migrateDialog.restartingServer') : t('models.migrate.offline')}

{migrationProgress.total > 0 && (
{migrationProgress.filename} {formatBytes(migrationProgress.current)} /{' '} {formatBytes(migrationProgress.total)}
)}
)}
); }