diff --git a/app/src/components/Generation/FloatingGenerateBox.tsx b/app/src/components/Generation/FloatingGenerateBox.tsx index 31b4d8ae..d4ab574d 100644 --- a/app/src/components/Generation/FloatingGenerateBox.tsx +++ b/app/src/components/Generation/FloatingGenerateBox.tsx @@ -316,7 +316,7 @@ export function FloatingGenerateBox({ - {isExpanded && form.watch('engine') !== 'luxtts' && ( + {isExpanded && form.watch('engine') === 'qwen' && ( { if (value === 'luxtts') { form.setValue('engine', 'luxtts'); + } else if (value === 'chatterbox') { + form.setValue('engine', 'chatterbox'); } else { const [, modelSize] = value.split(':'); form.setValue('engine', 'qwen'); @@ -434,6 +438,9 @@ export function FloatingGenerateBox({ LuxTTS + + Chatterbox + diff --git a/app/src/components/Generation/GenerationForm.tsx b/app/src/components/Generation/GenerationForm.tsx index 0ebabaa0..26fd13e3 100644 --- a/app/src/components/Generation/GenerationForm.tsx +++ b/app/src/components/Generation/GenerationForm.tsx @@ -76,7 +76,7 @@ export function GenerationForm() { )} /> - {form.watch('engine') !== 'luxtts' && ( + {form.watch('engine') === 'qwen' && ( { if (value === 'luxtts') { form.setValue('engine', 'luxtts'); + } else if (value === 'chatterbox') { + form.setValue('engine', 'chatterbox'); } else { const [, modelSize] = value.split(':'); form.setValue('engine', 'qwen'); @@ -128,12 +132,15 @@ export function GenerationForm() { Qwen3-TTS 1.7B Qwen3-TTS 0.6B LuxTTS + Chatterbox {form.watch('engine') === 'luxtts' ? 'Fast, English-focused' - : 'Multi-language, two sizes'} + : form.watch('engine') === 'chatterbox' + ? 'Multilingual, incl. Hebrew' + : 'Multi-language, two sizes'} diff --git a/app/src/components/ModelsTab/ModelsTab.tsx b/app/src/components/ModelsTab/ModelsTab.tsx index 3c6ebda7..4ef9f93b 100644 --- a/app/src/components/ModelsTab/ModelsTab.tsx +++ b/app/src/components/ModelsTab/ModelsTab.tsx @@ -2,7 +2,7 @@ import { ModelManagement } from '@/components/ServerSettings/ModelManagement'; export function ModelsTab() { return ( -
+
); diff --git a/app/src/components/ServerSettings/ModelManagement.tsx b/app/src/components/ServerSettings/ModelManagement.tsx index 56a2c3fb..4783a880 100644 --- a/app/src/components/ServerSettings/ModelManagement.tsx +++ b/app/src/components/ServerSettings/ModelManagement.tsx @@ -1,5 +1,21 @@ import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'; -import { ChevronDown, ChevronUp, Download, Loader2, RotateCcw, Trash2, X } from 'lucide-react'; +import { + ChevronDown, + ChevronRight, + ChevronUp, + CircleCheck, + CircleX, + Download, + ExternalLink, + HardDrive, + Heart, + Loader2, + RotateCcw, + Scale, + Trash2, + X, + Zap, +} from 'lucide-react'; import { useCallback, useState } from 'react'; import { AlertDialog, @@ -13,12 +29,50 @@ import { } from '@/components/ui/alert-dialog'; import { Badge } from '@/components/ui/badge'; import { Button } from '@/components/ui/button'; -import { Card, CardContent, CardDescription, CardHeader, CardTitle } from '@/components/ui/card'; +import { + Dialog, + DialogContent, + DialogDescription, + DialogHeader, + DialogTitle, +} from '@/components/ui/dialog'; import { useToast } from '@/components/ui/use-toast'; import { apiClient } from '@/lib/api/client'; -import type { ActiveDownloadTask } from '@/lib/api/types'; +import type { ActiveDownloadTask, HuggingFaceModelInfo, ModelStatus } from '@/lib/api/types'; import { useModelDownloadToast } from '@/lib/hooks/useModelDownloadToast'; +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(); +} + +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(' '); +} + export function ModelManagement() { const { toast } = useToast(); const queryClient = useQueryClient(); @@ -28,15 +82,17 @@ export function ModelManagement() { 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 () => { - console.log('[Query] Fetching model status'); const result = await apiClient.getModelStatus(); - console.log('[Query] Model status fetched:', result); return result; }, - refetchInterval: 5000, // Refresh every 5 seconds + refetchInterval: 5000, }); const { data: activeTasks } = useQuery({ @@ -45,19 +101,25 @@ export function ModelManagement() { refetchInterval: 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 - // Merge server errors with locally captured SSE errors const erroredDownloads = new Map(); if (activeTasks?.downloads) { for (const dl of activeTasks.downloads) { if (dl.status === 'error' && !dismissedErrors.has(dl.model_name)) { - // Prefer locally captured error (from SSE) over server error const localErr = localErrors.get(dl.model_name); erroredDownloads.set(dl.model_name, localErr ? { ...dl, error: localErr } : dl); } } } - // Also add locally captured errors that aren't in server response yet for (const [modelName, error] of localErrors) { if (!erroredDownloads.has(modelName) && !dismissedErrors.has(modelName)) { erroredDownloads.set(modelName, { @@ -71,9 +133,7 @@ export function ModelManagement() { const errorCount = erroredDownloads.size; - // Callbacks for download completion const handleDownloadComplete = useCallback(() => { - console.log('[ModelManagement] Download complete, clearing state'); setDownloadingModel(null); setDownloadingDisplayName(null); queryClient.invalidateQueries({ queryKey: ['modelStatus'] }); @@ -82,7 +142,6 @@ export function ModelManagement() { const handleDownloadError = useCallback( (error: string) => { - console.log('[ModelManagement] Download error, clearing state'); if (downloadingModel) { setLocalErrors((prev) => new Map(prev).set(downloadingModel, error)); setConsoleOpen(true); @@ -94,7 +153,6 @@ export function ModelManagement() { [queryClient, downloadingModel], ); - // Use progress toast hook for the downloading model useModelDownloadToast({ modelName: downloadingModel || '', displayName: downloadingDisplayName || '', @@ -111,36 +169,24 @@ export function ModelManagement() { } | null>(null); const handleDownload = async (modelName: string) => { - console.log('[Download] Button clicked for:', modelName, 'at', new Date().toISOString()); - // Clear any previous dismissal so fresh errors can appear setDismissedErrors((prev) => { const next = new Set(prev); next.delete(modelName); return next; }); - // Find display name const model = modelStatus?.models.find((m) => m.model_name === modelName); const displayName = model?.display_name || modelName; try { - // IMPORTANT: Call the API FIRST before setting state - // Setting state enables the SSE EventSource in useModelDownloadToast, - // which can block/delay the download fetch due to HTTP/1.1 connection limits - console.log('[Download] Calling download API for:', modelName); - const result = await apiClient.triggerModelDownload(modelName); - console.log('[Download] Download API responded:', result); + await apiClient.triggerModelDownload(modelName); - // NOW set state to enable SSE tracking (after download has started on backend) setDownloadingModel(modelName); setDownloadingDisplayName(displayName); - // Download initiated successfully - state will be cleared when SSE reports completion - // or by the polling interval detecting the model is downloaded queryClient.invalidateQueries({ queryKey: ['modelStatus'] }); queryClient.invalidateQueries({ queryKey: ['activeTasks'] }); } catch (error) { - console.error('[Download] Download failed:', error); setDownloadingModel(null); setDownloadingDisplayName(null); toast({ @@ -160,13 +206,11 @@ export function ModelManagement() { }); const handleCancel = (modelName: string) => { - // Snapshot previous state for rollback const prevDismissed = dismissedErrors; const prevLocalErrors = localErrors; const prevDownloadingModel = downloadingModel; const prevDownloadingDisplayName = downloadingDisplayName; - // Optimistically hide the error and suppress downloading state in UI setDismissedErrors((prev) => new Set(prev).add(modelName)); setLocalErrors((prev) => { const next = new Map(prev); @@ -180,7 +224,6 @@ export function ModelManagement() { cancelMutation.mutate(modelName, { onError: () => { - // Rollback optimistic updates on failure setDismissedErrors(prevDismissed); setLocalErrors(prevLocalErrors); setDownloadingModel(prevDownloadingModel); @@ -208,30 +251,22 @@ export function ModelManagement() { const deleteMutation = useMutation({ mutationFn: async (modelName: string) => { - console.log('[Delete] Deleting model:', modelName); const result = await apiClient.deleteModel(modelName); - console.log('[Delete] Model deleted successfully:', modelName); return result; }, - onSuccess: async (_data, _modelName) => { - console.log('[Delete] onSuccess - showing toast and invalidating queries'); + onSuccess: async () => { toast({ title: 'Model deleted', description: `${modelToDelete?.displayName || 'Model'} has been deleted successfully.`, }); setDeleteDialogOpen(false); setModelToDelete(null); - console.log('[Delete] Invalidating modelStatus query'); - await queryClient.invalidateQueries({ - queryKey: ['modelStatus'], - refetchType: 'all', - }); - console.log('[Delete] Explicitly refetching modelStatus query'); + setDetailOpen(false); + setSelectedModel(null); + await queryClient.invalidateQueries({ queryKey: ['modelStatus'], refetchType: 'all' }); await queryClient.refetchQueries({ queryKey: ['modelStatus'] }); - console.log('[Delete] Query refetched'); }, onError: (error: Error) => { - console.log('[Delete] onError:', error); toast({ title: 'Delete failed', description: error.message, @@ -241,185 +276,416 @@ export function ModelManagement() { }); const formatSize = (sizeMb?: number): string => { - if (!sizeMb) return 'Unknown'; + if (!sizeMb) return 'Unknown size'; 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 ttsModels = modelStatus?.models.filter((m) => m.model_name.startsWith('qwen-tts')) ?? []; + const otherTtsModels = + modelStatus?.models.filter( + (m) => m.model_name.startsWith('luxtts') || m.model_name.startsWith('chatterbox'), + ) ?? []; + const whisperModels = modelStatus?.models.filter((m) => m.model_name.startsWith('whisper')) ?? []; + + // Build sections + const sections: { label: string; models: ModelStatus[] }[] = [ + { label: 'Voice Generation', models: ttsModels }, + ...(otherTtsModels.length > 0 ? [{ label: 'Other Voice Models', models: otherTtsModels }] : []), + { label: 'Transcription', models: whisperModels }, + ]; + + // 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((t) => t.startsWith('license:'))?.replace('license:', ''); + return ( - - - Model Management - +
+ {/* Header */} +
+

Models

+

Download and manage AI models for voice generation and transcription - - - - {isLoading ? ( -

- -
- ) : modelStatus ? ( -
- {/* TTS Models */} -
-

- Voice Generation Models -

-
- {modelStatus.models - .filter((m) => m.model_name.startsWith('qwen-tts')) - .map((model) => ( - +
+ + {/* Model list */} + {isLoading ? ( +
+ +
+ ) : modelStatus ? ( +
+ {sections.map((section) => ( +
+

+ {section.label} +

+
+ {section.models.map((model) => { + const { isDownloading, hasError } = getModelState(model); + return ( +
-
- - {/* LuxTTS Models */} - {modelStatus.models.some((m) => m.model_name.startsWith('luxtts')) && ( -
-

LuxTTS Models

-
- {modelStatus.models - .filter((m) => m.model_name.startsWith('luxtts')) - .map((model) => ( - handleDownload(model.model_name)} - onDelete={() => { - setModelToDelete({ - name: model.model_name, - displayName: model.display_name, - sizeMb: model.size_mb, - }); - setDeleteDialogOpen(true); - }} - isDownloading={downloadingModel === model.model_name} - formatSize={formatSize} - /> - ))} -
-
- )} - - {/* Whisper Models */} -
-

- Transcription Models -

-
- {modelStatus.models - .filter((m) => m.model_name.startsWith('whisper')) - .map((model) => ( - handleDownload(model.model_name)} - onDelete={() => { - setModelToDelete({ - name: model.model_name, - displayName: model.display_name, - sizeMb: model.size_mb, - }); - setDeleteDialogOpen(true); - }} - onCancel={() => handleCancel(model.model_name)} - isDownloading={downloadingModel === model.model_name} - isCancelling={ - cancelMutation.isPending && cancelMutation.variables === model.model_name - } - isDismissed={dismissedErrors.has(model.model_name)} - erroredDownload={erroredDownloads.get(model.model_name)} - formatSize={formatSize} - /> - ))} -
-
- - {/* Console Panel */} - {errorCount > 0 && ( -
-
- - -
- {consoleOpen && ( -
- {Array.from(erroredDownloads.entries()).map(([modelName, dl]) => ( -
- [error]{' '} - {modelName} - {dl.error ? ( - <> - {': '} - - {dl.error} - - + type="button" + onClick={() => openModelDetail(model)} + className="w-full flex items-center gap-3 px-3 py-2.5 text-left hover:bg-muted/50 transition-colors group" + > + {/* Status indicator */} +
+ {hasError ? ( + + ) : isDownloading ? ( + + ) : model.loaded ? ( + + ) : model.downloaded ? ( + ) : ( - <> - {': '} - - No error details available. Try downloading again. - - + )} -
- started at {new Date(dl.started_at).toLocaleString()} -
- ))} + + {/* Name + meta */} +
+ {model.display_name} +
+ + {/* Right side info */} +
+ {isDownloading && ( + Downloading... + )} + {hasError && ( + + Error + + )} + {model.loaded && ( + + Active + + )} + {model.downloaded && !model.loaded && !isDownloading && !hasError && ( + + {formatSize(model.size_mb)} + + )} + {!model.downloaded && !isDownloading && !hasError && ( + Not downloaded + )} + +
+ + ); + })} +
+
+ ))} + + {/* Error console */} + {errorCount > 0 && ( +
+
+ + +
+ {consoleOpen && ( +
+ {Array.from(erroredDownloads.entries()).map(([modelName, dl]) => ( +
+ [error]{' '} + {modelName} + {dl.error ? ( + <> + {': '} + + {dl.error} + + + ) : ( + <> + {': '} + + No error details available. Try downloading again. + + + )} +
+ started at {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 && ( + + + Loaded in memory + + )} + {freshSelectedModel.downloaded && !freshSelectedModel.loaded && ( + + + Downloaded + + )} + {selectedState?.isDownloading && ( + + + Downloading + + )} + {selectedState?.hasError && ( + + + Error + + )} + {!freshSelectedModel.downloaded && + !selectedState?.isDownloading && + !selectedState?.hasError && ( + + Not downloaded + + )} +
+ + {/* HuggingFace model card info */} + {hfLoading && freshSelectedModel.hf_repo_id && ( +
+ + Loading model info...
)} + + {hfModelInfo && ( +
+ {/* Stats row */} +
+ + + {formatDownloads(hfModelInfo.downloads)} + + + + {formatDownloads(hfModelInfo.likes)} + + {license && ( + + + {formatLicense(license)} + + )} +
+ + {/* Pipeline tag + author */} +
+ {hfModelInfo.pipeline_tag && ( + + {formatPipelineTag(hfModelInfo.pipeline_tag)} + + )} + {hfModelInfo.library_name && ( + + {hfModelInfo.library_name} + + )} + {hfModelInfo.author && ( + + by {hfModelInfo.author} + + )} +
+ + {/* Languages */} + {hfModelInfo.cardData?.language && hfModelInfo.cardData.language.length > 0 && ( +
+ + {hfModelInfo.cardData.language.length > 10 + ? `${hfModelInfo.cardData.language.length} languages supported` + : `Languages: ${hfModelInfo.cardData.language.join(', ')}`} + +
+ )} +
+ )} + + {/* Disk size */} + {freshSelectedModel.downloaded && freshSelectedModel.size_mb && ( +
+ + {formatSize(freshSelectedModel.size_mb)} on disk +
+ )} + + {/* Error detail */} + {selectedError?.error && ( +
+ {selectedError.error} +
+ )} + + {/* Actions */} +
+ {selectedState?.hasError ? ( + <> + + + + ) : selectedState?.isDownloading ? ( + <> + + + + ) : freshSelectedModel.downloaded ? ( + + ) : ( + + )} +
- )} -
- ) : null} - + + )} + + {/* Delete Confirmation Dialog */} @@ -460,126 +726,6 @@ export function ModelManagement() { - - ); -} - -interface ModelItemProps { - model: { - model_name: string; - display_name: string; - downloaded: boolean; - downloading?: boolean; // From server - true if download in progress - size_mb?: number; - loaded: boolean; - }; - onDownload: () => void; - onDelete: () => void; - onCancel: () => void; - isDownloading: boolean; // Local state - true if user just clicked download - isCancelling: boolean; - isDismissed: boolean; - erroredDownload?: ActiveDownloadTask; - formatSize: (sizeMb?: number) => string; -} - -function ModelItem({ - model, - onDownload, - onDelete, - onCancel, - isDownloading, - isCancelling, - isDismissed, - erroredDownload, - formatSize, -}: ModelItemProps) { - // Use server's downloading state OR local state (for immediate feedback before server updates) - // Suppress downloading if user just dismissed/cancelled this model - const showDownloading = (model.downloading || isDownloading) && !erroredDownload && !isDismissed; - - return ( -
-
-
- {model.display_name} - {model.loaded && ( - - Loaded - - )} - {model.downloaded && !model.loaded && !showDownloading && !erroredDownload && ( - - Downloaded - - )} - {erroredDownload && ( - - Error - - )} -
- {model.downloaded && model.size_mb && !showDownloading && !erroredDownload && ( -
- Size: {formatSize(model.size_mb)} -
- )} -
-
- {erroredDownload ? ( -
- - -
- ) : model.downloaded && !showDownloading ? ( -
-
- Ready -
- -
- ) : showDownloading ? ( -
- - -
- ) : ( - - )} -
); } diff --git a/app/src/lib/api/types.ts b/app/src/lib/api/types.ts index 4041a318..0da89da5 100644 --- a/app/src/lib/api/types.ts +++ b/app/src/lib/api/types.ts @@ -34,7 +34,7 @@ export interface GenerationRequest { language: LanguageCode; seed?: number; model_size?: '1.7B' | '0.6B'; - engine?: 'qwen' | 'luxtts'; + engine?: 'qwen' | 'luxtts' | 'chatterbox'; instruct?: string; } @@ -119,12 +119,29 @@ export interface ModelProgress { export interface ModelStatus { model_name: string; display_name: string; + hf_repo_id?: string; // HuggingFace repository ID downloaded: boolean; downloading: boolean; // True if download is in progress size_mb?: number; loaded: boolean; } +export interface HuggingFaceModelInfo { + id: string; + author: string; + lastModified: string; + pipeline_tag?: string; + library_name?: string; + downloads: number; + likes: number; + tags: string[]; + cardData?: { + license?: string; + language?: string[]; + pipeline_tag?: string; + }; +} + export interface ModelStatusListResponse { models: ModelStatus[]; } diff --git a/app/src/lib/constants/languages.ts b/app/src/lib/constants/languages.ts index 9ffc396f..9c85aafd 100644 --- a/app/src/lib/constants/languages.ts +++ b/app/src/lib/constants/languages.ts @@ -1,6 +1,6 @@ /** - * Supported languages for Qwen3-TTS - * Based on: https://github.com/QwenLM/Qwen3-TTS + * Supported languages for voice generation. + * Most languages use Qwen3-TTS; Hebrew uses Chatterbox TTS. */ export const SUPPORTED_LANGUAGES = { @@ -14,6 +14,7 @@ export const SUPPORTED_LANGUAGES = { pt: 'Portuguese', es: 'Spanish', it: 'Italian', + he: 'Hebrew', } as const; export type LanguageCode = keyof typeof SUPPORTED_LANGUAGES; diff --git a/app/src/lib/hooks/useGenerationForm.ts b/app/src/lib/hooks/useGenerationForm.ts index 1a44b24f..ec5b9d6a 100644 --- a/app/src/lib/hooks/useGenerationForm.ts +++ b/app/src/lib/hooks/useGenerationForm.ts @@ -16,7 +16,7 @@ const generationSchema = z.object({ seed: z.number().int().optional(), modelSize: z.enum(['1.7B', '0.6B']).optional(), instruct: z.string().max(500).optional(), - engine: z.enum(['qwen', 'luxtts']).optional(), + engine: z.enum(['qwen', 'luxtts', 'chatterbox']).optional(), }); export type GenerationFormValues = z.infer; @@ -70,13 +70,20 @@ export function useGenerationForm(options: UseGenerationFormOptions = {}) { setIsGenerating(true); const engine = data.engine || 'qwen'; - const modelName = engine === 'luxtts' ? 'luxtts' : `qwen-tts-${data.modelSize}`; + const modelName = + engine === 'luxtts' + ? 'luxtts' + : engine === 'chatterbox' + ? 'chatterbox-tts' + : `qwen-tts-${data.modelSize}`; const displayName = engine === 'luxtts' ? 'LuxTTS' - : data.modelSize === '1.7B' - ? 'Qwen TTS 1.7B' - : 'Qwen TTS 0.6B'; + : engine === 'chatterbox' + ? 'Chatterbox TTS' + : data.modelSize === '1.7B' + ? 'Qwen TTS 1.7B' + : 'Qwen TTS 0.6B'; try { const modelStatus = await apiClient.getModelStatus(); @@ -90,14 +97,15 @@ export function useGenerationForm(options: UseGenerationFormOptions = {}) { console.error('Failed to check model status:', error); } + const isQwen = engine === 'qwen'; const result = await generation.mutateAsync({ profile_id: selectedProfileId, text: data.text, language: data.language, seed: data.seed, - model_size: engine === 'luxtts' ? undefined : data.modelSize, + model_size: isQwen ? data.modelSize : undefined, engine, - instruct: engine === 'luxtts' ? undefined : data.instruct || undefined, + instruct: isQwen ? data.instruct || undefined : undefined, }); toast({ diff --git a/backend/backends/__init__.py b/backend/backends/__init__.py index 8f4dce0c..a7b4d54c 100644 --- a/backend/backends/__init__.py +++ b/backend/backends/__init__.py @@ -121,6 +121,7 @@ _stt_backend: Optional[STTBackend] = None TTS_ENGINES = { "qwen": "Qwen TTS", "luxtts": "LuxTTS", + "chatterbox": "Chatterbox TTS", } @@ -167,6 +168,9 @@ def get_tts_backend_for_engine(engine: str) -> TTSBackend: elif engine == "luxtts": from .luxtts_backend import LuxTTSBackend backend = LuxTTSBackend() + elif engine == "chatterbox": + from .chatterbox_backend import ChatterboxTTSBackend + backend = ChatterboxTTSBackend() else: raise ValueError(f"Unknown TTS engine: {engine}. Supported: {list(TTS_ENGINES.keys())}") diff --git a/backend/backends/chatterbox_backend.py b/backend/backends/chatterbox_backend.py new file mode 100644 index 00000000..84eb8365 --- /dev/null +++ b/backend/backends/chatterbox_backend.py @@ -0,0 +1,318 @@ +""" +Chatterbox TTS backend implementation. + +Wraps ChatterboxMultilingualTTS from chatterbox-tts for zero-shot +voice cloning. Supports 23 languages including Hebrew. Forces CPU +on macOS due to known MPS tensor issues. +""" + +import asyncio +import logging +import platform +import threading +from pathlib import Path +from typing import ClassVar, List, Optional, Tuple + +import numpy as np + +from . import TTSBackend +from ..utils.audio import normalize_audio, load_audio +from ..utils.progress import get_progress_manager +from ..utils.hf_progress import HFProgressTracker, create_hf_progress_callback +from ..utils.tasks import get_task_manager + +logger = logging.getLogger(__name__) + +CHATTERBOX_HF_REPO = "ResembleAI/chatterbox" + +# Files that must be present for the multilingual model +_MTL_WEIGHT_FILES = [ + "t3_mtl23ls_v2.safetensors", + "s3gen.pt", + "ve.pt", +] + + +class ChatterboxTTSBackend: + """Chatterbox Multilingual TTS backend for voice cloning.""" + + # Class-level lock for torch.load monkey-patching + _load_lock: ClassVar[threading.Lock] = threading.Lock() + + def __init__(self): + self.model = None + self.model_size = "default" + self._device = None + self._model_load_lock = asyncio.Lock() + + def _get_device(self) -> str: + """Get the best available device. Forces CPU on macOS (MPS issue).""" + if platform.system() == "Darwin": + return "cpu" + try: + import torch + + if torch.cuda.is_available(): + return "cuda" + except ImportError: + pass + return "cpu" + + def is_loaded(self) -> bool: + return self.model is not None + + def _get_model_path(self, model_size: str = "default") -> str: + return CHATTERBOX_HF_REPO + + def _is_model_cached(self, model_size: str = "default") -> bool: + """Check if the Chatterbox multilingual model is cached locally.""" + try: + from huggingface_hub import constants as hf_constants + + repo_cache = Path(hf_constants.HF_HUB_CACHE) / ( + "models--" + CHATTERBOX_HF_REPO.replace("/", "--") + ) + + if not repo_cache.exists(): + return False + + blobs_dir = repo_cache / "blobs" + if blobs_dir.exists() and any(blobs_dir.glob("*.incomplete")): + return False + + # Check for multilingual weight files + snapshots_dir = repo_cache / "snapshots" + if snapshots_dir.exists(): + for fname in _MTL_WEIGHT_FILES: + if not any(snapshots_dir.rglob(fname)): + return False + return True + + return False + except Exception as e: + logger.warning(f"Error checking Chatterbox cache: {e}") + return False + + async def load_model(self, model_size: str = "default") -> None: + """Load the Chatterbox multilingual model.""" + if self.model is not None: + return + async with self._model_load_lock: + if self.model is not None: + return + await asyncio.to_thread(self._load_model_sync) + + def _load_model_sync(self): + """Synchronous model loading.""" + progress_manager = get_progress_manager() + task_manager = get_task_manager() + model_name = "chatterbox-tts" + + is_cached = self._is_model_cached() + + try: + progress_callback = create_hf_progress_callback(model_name, progress_manager) + tracker = HFProgressTracker(progress_callback, filter_non_downloads=is_cached) + + if not is_cached: + task_manager.start_download(model_name) + progress_manager.update_progress( + model_name=model_name, + current=0, + total=0, + filename="Downloading Chatterbox model...", + status="downloading", + ) + + with tracker.patch_download(): + device = self._get_device() + self._device = device + + logger.info(f"Loading Chatterbox Multilingual TTS on {device}...") + + import torch + from chatterbox.mtl_tts import ChatterboxMultilingualTTS + + # Monkey-patch torch.load for CPU loading. The model's .pt files + # were saved on CUDA; from_pretrained() doesn't pass map_location + # so loading on CPU fails without this. + if device == "cpu": + _orig_torch_load = torch.load + + def _patched_load(*args, **kwargs): + kwargs.setdefault("map_location", "cpu") + return _orig_torch_load(*args, **kwargs) + + with ChatterboxTTSBackend._load_lock: + torch.load = _patched_load + try: + self.model = ChatterboxMultilingualTTS.from_pretrained( + device=device, + ) + finally: + torch.load = _orig_torch_load + else: + self.model = ChatterboxMultilingualTTS.from_pretrained( + device=device, + ) + + # Fix: transformers >= 4.36 defaults LlamaModel to sdpa attention + # which doesn't support output_attentions=True (needed by + # Chatterbox's AlignmentStreamAnalyzer). Force eager attention. + t3_tfmr = self.model.t3.tfmr + if hasattr(t3_tfmr, "config") and hasattr( + t3_tfmr.config, "_attn_implementation" + ): + t3_tfmr.config._attn_implementation = "eager" + for layer in getattr(t3_tfmr, "layers", []): + if hasattr(layer, "self_attn"): + layer.self_attn._attn_implementation = "eager" + + if not is_cached: + progress_manager.mark_complete(model_name) + task_manager.complete_download(model_name) + + logger.info("Chatterbox Multilingual TTS loaded successfully") + + except ImportError as e: + logger.error( + "chatterbox-tts package not found. " + "Install with: pip install chatterbox-tts" + ) + progress_manager.mark_error(model_name, str(e)) + task_manager.error_download(model_name, str(e)) + raise + except Exception as e: + logger.error(f"Failed to load Chatterbox: {e}") + progress_manager.mark_error(model_name, str(e)) + task_manager.error_download(model_name, str(e)) + raise + + def unload_model(self) -> None: + """Unload model to free memory.""" + if self.model is not None: + device = self._device + del self.model + self.model = None + self._device = None + if device == "cuda": + import torch + + torch.cuda.empty_cache() + logger.info("Chatterbox unloaded") + + async def create_voice_prompt( + self, + audio_path: str, + reference_text: str, + use_cache: bool = True, + ) -> Tuple[dict, bool]: + """ + Create voice prompt from reference audio. + + Chatterbox processes reference audio at generation time, so the + prompt just stores the file path. The actual audio is loaded by + model.generate() via audio_prompt_path. + """ + voice_prompt = { + "ref_audio": str(audio_path), + "ref_text": reference_text, + } + return voice_prompt, False + + async def combine_voice_prompts( + self, + audio_paths: List[str], + reference_texts: List[str], + ) -> Tuple[np.ndarray, str]: + """Combine multiple reference samples.""" + combined_audio = [] + for path in audio_paths: + audio, _sr = load_audio(path) + audio = normalize_audio(audio) + combined_audio.append(audio) + + mixed = np.concatenate(combined_audio) + mixed = normalize_audio(mixed) + combined_text = " ".join(reference_texts) + return mixed, combined_text + + # Per-language generation defaults. Lower temp + higher cfg = clearer speech. + _LANG_DEFAULTS: ClassVar[dict] = { + "he": { + "exaggeration": 0.4, + "cfg_weight": 0.7, + "temperature": 0.65, + "repetition_penalty": 2.5, + }, + } + _GLOBAL_DEFAULTS: ClassVar[dict] = { + "exaggeration": 0.5, + "cfg_weight": 0.5, + "temperature": 0.8, + "repetition_penalty": 2.0, + } + + async def generate( + self, + text: str, + voice_prompt: dict, + language: str = "en", + seed: Optional[int] = None, + instruct: Optional[str] = None, + ) -> Tuple[np.ndarray, int]: + """ + Generate audio using Chatterbox Multilingual TTS. + + Args: + text: Text to synthesize + voice_prompt: Dict with ref_audio path + language: BCP-47 language code + seed: Random seed for reproducibility + instruct: Unused (protocol compatibility) + + Returns: + Tuple of (audio_array, sample_rate) + """ + await self.load_model() + + ref_audio = voice_prompt.get("ref_audio") + if ref_audio and not Path(ref_audio).exists(): + logger.warning(f"Reference audio not found: {ref_audio}") + ref_audio = None + + # Merge language-specific defaults with global defaults + lang_defaults = self._LANG_DEFAULTS.get(language, self._GLOBAL_DEFAULTS) + + def _generate_sync(): + import torch + + if seed is not None: + torch.manual_seed(seed) + + logger.info(f"[Chatterbox] Generating: lang={language}") + + wav = self.model.generate( + text, + language_id=language, + audio_prompt_path=ref_audio, + exaggeration=lang_defaults["exaggeration"], + cfg_weight=lang_defaults["cfg_weight"], + temperature=lang_defaults["temperature"], + repetition_penalty=lang_defaults["repetition_penalty"], + ) + + # Convert tensor -> numpy + if isinstance(wav, torch.Tensor): + audio = wav.squeeze().cpu().numpy().astype(np.float32) + else: + audio = np.asarray(wav, dtype=np.float32) + + sample_rate = ( + getattr(self.model, "sr", None) + or getattr(self.model, "sample_rate", 24000) + ) + + return audio, sample_rate + + return await asyncio.to_thread(_generate_sync) diff --git a/backend/main.py b/backend/main.py index d4dd87c5..7eeabe40 100644 --- a/backend/main.py +++ b/backend/main.py @@ -676,6 +676,29 @@ async def generate_speech( ) await tts_model.load_model() + elif engine == "chatterbox": + if not tts_model._is_model_cached(): + model_name = "chatterbox-tts" + + async def download_chatterbox_background(): + try: + await tts_model.load_model() + except Exception as e: + task_manager.error_download(model_name, str(e)) + + task_manager.start_download(model_name) + asyncio.create_task(download_chatterbox_background()) + + raise HTTPException( + status_code=202, + detail={ + "message": "Chatterbox model is being downloaded. Please wait and try again.", + "model_name": model_name, + "downloading": True, + }, + ) + + await tts_model.load_model() # Create voice prompt from profile voice_prompt = await profiles.create_voice_prompt_for_profile( @@ -693,6 +716,11 @@ async def generate_speech( data.instruct, ) + # Trim trailing silence/hallucination for Chatterbox output + if engine == "chatterbox": + from .utils.audio import trim_tts_output + audio = trim_tts_output(audio, sample_rate) + # Calculate duration duration = len(audio) / sample_rate @@ -763,6 +791,13 @@ async def stream_speech( detail="LuxTTS model is not downloaded yet. Use /generate to trigger a download.", ) await tts_model.load_model() + elif engine == "chatterbox": + if not tts_model._is_model_cached(): + raise HTTPException( + status_code=400, + detail="Chatterbox model is not downloaded yet. Use /generate to trigger a download.", + ) + await tts_model.load_model() voice_prompt = await profiles.create_voice_prompt_for_profile( data.profile_id, db, engine=engine, @@ -776,6 +811,11 @@ async def stream_speech( data.instruct, ) + # Trim trailing silence/hallucination for Chatterbox output + if engine == "chatterbox": + from .utils.audio import trim_tts_output + audio = trim_tts_output(audio, sample_rate) + wav_bytes = tts.audio_to_wav_bytes(audio, sample_rate) async def _wav_stream(): @@ -1384,6 +1424,15 @@ async def get_model_status(): except Exception: return False + # Check if Chatterbox backend is loaded + def check_chatterbox_loaded(): + try: + from .backends import get_tts_backend_for_engine + backend = get_tts_backend_for_engine("chatterbox") + return backend.is_loaded() + except Exception: + return False + model_configs = [ { "model_name": "qwen-tts-1.7B", @@ -1406,6 +1455,13 @@ async def get_model_status(): "model_size": "default", "check_loaded": check_luxtts_loaded, }, + { + "model_name": "chatterbox-tts", + "display_name": "Chatterbox TTS (Multilingual)", + "hf_repo_id": "ResembleAI/chatterbox", + "model_size": "default", + "check_loaded": check_chatterbox_loaded, + }, { "model_name": "whisper-base", "display_name": "Whisper Base", @@ -1557,6 +1613,7 @@ async def get_model_status(): statuses.append(models.ModelStatus( model_name=config["model_name"], display_name=config["display_name"], + hf_repo_id=config["hf_repo_id"], downloaded=downloaded, downloading=is_downloading, size_mb=size_mb, @@ -1575,6 +1632,7 @@ async def get_model_status(): statuses.append(models.ModelStatus( model_name=config["model_name"], display_name=config["display_name"], + hf_repo_id=config["hf_repo_id"], downloaded=False, # Assume not downloaded if check failed downloading=is_downloading, size_mb=None, @@ -1606,6 +1664,10 @@ async def trigger_model_download(request: models.ModelDownloadRequest): "model_size": "default", "load_func": lambda: get_tts_backend_for_engine("luxtts").load_model(), }, + "chatterbox-tts": { + "model_size": "default", + "load_func": lambda: get_tts_backend_for_engine("chatterbox").load_model(), + }, "whisper-base": { "model_size": "base", "load_func": lambda: transcribe.get_whisper_model().load_model("base"), @@ -1723,6 +1785,11 @@ async def delete_model(model_name: str): "model_size": "default", "model_type": "luxtts", }, + "chatterbox-tts": { + "hf_repo_id": "ResembleAI/chatterbox", + "model_size": "default", + "model_type": "chatterbox", + }, "whisper-base": { "hf_repo_id": "openai/whisper-base", "model_size": "base", @@ -1762,6 +1829,11 @@ async def delete_model(model_name: str): luxtts = get_tts_backend_for_engine("luxtts") if luxtts.is_loaded(): luxtts.unload_model() + elif config["model_type"] == "chatterbox": + from .backends import get_tts_backend_for_engine + chatterbox = get_tts_backend_for_engine("chatterbox") + if chatterbox.is_loaded(): + chatterbox.unload_model() elif config["model_type"] == "whisper": whisper_model = transcribe.get_whisper_model() if whisper_model.is_loaded() and whisper_model.model_size == config["model_size"]: diff --git a/backend/models.py b/backend/models.py index 9ecf4510..9b945adb 100644 --- a/backend/models.py +++ b/backend/models.py @@ -11,7 +11,7 @@ class VoiceProfileCreate(BaseModel): """Request model for creating a voice profile.""" name: str = Field(..., min_length=1, max_length=100) description: Optional[str] = Field(None, max_length=500) - language: str = Field(default="en", pattern="^(zh|en|ja|ko|de|fr|ru|pt|es|it)$") + language: str = Field(default="en", pattern="^(zh|en|ja|ko|de|fr|ru|pt|es|it|he)$") class VoiceProfileResponse(BaseModel): @@ -53,11 +53,11 @@ class GenerationRequest(BaseModel): """Request model for voice generation.""" profile_id: str text: str = Field(..., min_length=1, max_length=5000) - language: str = Field(default="en", pattern="^(zh|en|ja|ko|de|fr|ru|pt|es|it)$") + language: str = Field(default="en", pattern="^(zh|en|ja|ko|de|fr|ru|pt|es|it|he)$") seed: Optional[int] = Field(None, ge=0) model_size: Optional[str] = Field(default="1.7B", pattern="^(1\\.7B|0\\.6B)$") instruct: Optional[str] = Field(None, max_length=500) - engine: Optional[str] = Field(default="qwen", pattern="^(qwen|luxtts)$") + engine: Optional[str] = Field(default="qwen", pattern="^(qwen|luxtts|chatterbox)$") class GenerationResponse(BaseModel): @@ -135,6 +135,7 @@ class ModelStatus(BaseModel): """Response model for model status.""" model_name: str display_name: str + hf_repo_id: Optional[str] = None # HuggingFace repository ID downloaded: bool downloading: bool = False # True if download is in progress size_mb: Optional[float] = None diff --git a/backend/requirements.txt b/backend/requirements.txt index e57ddd94..53e9e6d3 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -21,6 +21,9 @@ qwen-tts>=0.0.5 linacodec @ git+https://github.com/ysharma3501/LinaCodec.git Zipvoice @ git+https://github.com/ysharma3501/LuxTTS.git +# Chatterbox TTS (multilingual voice cloning, includes Hebrew) +chatterbox-tts>=0.1.0 + # Audio processing librosa>=0.10.0 soundfile>=0.12.0 diff --git a/backend/utils/audio.py b/backend/utils/audio.py index 302dff25..20709d92 100644 --- a/backend/utils/audio.py +++ b/backend/utils/audio.py @@ -80,6 +80,95 @@ def save_audio( sf.write(path, audio, sample_rate) +def trim_tts_output( + audio: np.ndarray, + sample_rate: int = 24000, + frame_ms: int = 20, + silence_threshold_db: float = -40.0, + min_silence_ms: int = 200, + max_internal_silence_ms: int = 1000, + fade_ms: int = 30, +) -> np.ndarray: + """ + Trim trailing silence and post-silence hallucination from TTS output. + + Chatterbox sometimes produces ``[speech][silence][hallucinated noise]``. + This detects internal silence gaps longer than *max_internal_silence_ms* + and cuts the audio at that boundary, then trims trailing silence and + applies a short cosine fade-out. + + Args: + audio: Input audio array (mono float32) + sample_rate: Sample rate in Hz + frame_ms: Frame size for RMS energy calculation + silence_threshold_db: dB threshold below which a frame is silence + min_silence_ms: Minimum trailing silence to keep + max_internal_silence_ms: Cut after any silence gap longer than this + fade_ms: Cosine fade-out duration in ms + + Returns: + Trimmed audio array + """ + frame_len = int(sample_rate * frame_ms / 1000) + if frame_len == 0 or len(audio) < frame_len: + return audio + + n_frames = len(audio) // frame_len + threshold_linear = 10 ** (silence_threshold_db / 20) + + # Compute per-frame RMS + rms = np.array( + [ + np.sqrt(np.mean(audio[i * frame_len : (i + 1) * frame_len] ** 2)) + for i in range(n_frames) + ] + ) + is_speech = rms >= threshold_linear + + # Find first speech frame + first_speech = 0 + for i, s in enumerate(is_speech): + if s: + first_speech = max(0, i - 1) # keep 1 frame padding + break + + # Walk forward from first speech; cut at long internal silence gaps + max_silence_frames = int(max_internal_silence_ms / frame_ms) + consecutive_silence = 0 + cut_frame = n_frames + + for i in range(first_speech, n_frames): + if is_speech[i]: + consecutive_silence = 0 + else: + consecutive_silence += 1 + if consecutive_silence >= max_silence_frames: + cut_frame = i - consecutive_silence + 1 + break + + # Trim trailing silence from the cut point + min_silence_frames = int(min_silence_ms / frame_ms) + end_frame = cut_frame + while end_frame > first_speech and not is_speech[end_frame - 1]: + end_frame -= 1 + # Keep a short tail + end_frame = min(end_frame + min_silence_frames, cut_frame) + + # Convert frames back to samples + start_sample = first_speech * frame_len + end_sample = min(end_frame * frame_len, len(audio)) + + trimmed = audio[start_sample:end_sample].copy() + + # Cosine fade-out + fade_samples = int(sample_rate * fade_ms / 1000) + if fade_samples > 0 and len(trimmed) > fade_samples: + fade = np.cos(np.linspace(0, np.pi / 2, fade_samples)) ** 2 + trimmed[-fade_samples:] *= fade + + return trimmed + + def validate_reference_audio( audio_path: str, min_duration: float = 2.0, diff --git a/scripts/test_download_progress.py b/scripts/test_download_progress.py new file mode 100644 index 00000000..a0fbe551 --- /dev/null +++ b/scripts/test_download_progress.py @@ -0,0 +1,382 @@ +#!/usr/bin/env python3 +""" +Test script to observe exactly how HuggingFace reports download progress +for each TTS model. Doesn't load models — just downloads and tracks tqdm. + +Usage: + backend/venv/bin/python scripts/test_download_progress.py qwen + backend/venv/bin/python scripts/test_download_progress.py luxtts + backend/venv/bin/python scripts/test_download_progress.py chatterbox + +Add --delete to clear cache first and force a real download: + backend/venv/bin/python scripts/test_download_progress.py chatterbox --delete +""" + +import os +import shutil +import sys +import time +import threading +from pathlib import Path +from contextlib import contextmanager + +# ─── Configuration ──────────────────────────────────────────────────────────── + +MODELS = { + "qwen": { + "repo_id": "Qwen/Qwen3-TTS-12Hz-1.7B-Base", + "method": "from_pretrained", + "description": "Qwen TTS 1.7B (uses transformers from_pretrained)", + }, + "luxtts": { + "repo_id": "YatharthS/LuxTTS", + "method": "snapshot_download", + "description": "LuxTTS (uses snapshot_download)", + }, + "chatterbox": { + "repo_id": "ResembleAI/chatterbox", + "method": "snapshot_download", + "allow_patterns": [ + "ve.pt", + "t3_mtl23ls_v2.safetensors", + "s3gen.pt", + "grapheme_mtl_merged_expanded_v1.json", + "conds.pt", + "Cangjie5_TC.json", + ], + "description": "Chatterbox Multilingual (uses snapshot_download with allow_patterns)", + }, +} + + +# ─── Progress tracking (mirrors our HFProgressTracker) ──────────────────────── + +class ProgressSpy: + """Intercepts tqdm to see exactly what HF reports.""" + + def __init__(self): + self._lock = threading.Lock() + self.events = [] # List of dicts: {time, type, ...} + self._original_tqdm_class = None + self._original_tqdm_auto = None + self._patched_modules = {} + self._hf_tqdm_original_update = None + self._start_time = None + + def _elapsed(self): + return time.time() - self._start_time if self._start_time else 0 + + def _log(self, event_type, **kwargs): + entry = {"time": f"{self._elapsed():.1f}s", "type": event_type, **kwargs} + self.events.append(entry) + + # Live print + parts = [f"[{entry['time']:>7s}] {event_type:>10s}"] + for k, v in kwargs.items(): + if k in ("current", "total") and isinstance(v, (int, float)) and v > 1_000_000: + parts.append(f"{k}={v / 1_000_000:.1f}MB") + else: + parts.append(f"{k}={v}") + print(" ".join(parts), flush=True) + + def _create_tracked_tqdm_class(self): + spy = self + original_tqdm = self._original_tqdm_class + + class SpyTqdm(original_tqdm): + def __init__(self, *args, **kwargs): + desc = kwargs.get("desc", "") + if not desc and args: + first_arg = args[0] + if isinstance(first_arg, str): + desc = first_arg + + filename = "" + if desc: + if ":" in desc: + filename = desc.split(":")[0].strip() + else: + filename = desc.strip() + + # Filter out non-standard kwargs + tqdm_kwargs = { + 'iterable', 'desc', 'total', 'leave', 'file', 'ncols', + 'mininterval', 'maxinterval', 'miniters', 'ascii', 'disable', + 'unit', 'unit_scale', 'dynamic_ncols', 'smoothing', + 'bar_format', 'initial', 'position', 'postfix', + 'unit_divisor', 'write_bytes', 'lock_args', 'nrows', + 'colour', 'color', 'delay', 'gui', 'disable_default', 'pos', + } + filtered_kwargs = {k: v for k, v in kwargs.items() if k in tqdm_kwargs} + + try: + super().__init__(*args, **filtered_kwargs) + except TypeError: + super().__init__(*args, **kwargs) + + self._spy_filename = filename or "unknown" + total = getattr(self, "total", None) + + spy._log( + "INIT", + filename=self._spy_filename, + total=total or 0, + unit=kwargs.get("unit", "?"), + unit_scale=kwargs.get("unit_scale", False), + disable=kwargs.get("disable", False), + ) + + def update(self, n=1): + result = super().update(n) + + current = getattr(self, "n", 0) + total = getattr(self, "total", 0) + filename = self._spy_filename + + spy._log( + "UPDATE", + filename=filename, + n=n, + current=current, + total=total or 0, + pct=f"{100 * current / total:.1f}%" if total else "?", + ) + + return result + + def close(self): + spy._log("CLOSE", filename=self._spy_filename) + return super().close() + + return SpyTqdm + + @contextmanager + def patch(self): + """Context manager that patches tqdm globally — same as HFProgressTracker.""" + self._start_time = time.time() + + try: + import tqdm as tqdm_module + self._original_tqdm_class = tqdm_module.tqdm + except ImportError: + yield + return + + tracked_tqdm = self._create_tracked_tqdm_class() + + # Patch tqdm.tqdm + tqdm_module.tqdm = tracked_tqdm + + # Patch tqdm.auto.tqdm + self._original_tqdm_auto = None + if hasattr(tqdm_module, "auto") and hasattr(tqdm_module.auto, "tqdm"): + self._original_tqdm_auto = tqdm_module.auto.tqdm + tqdm_module.auto.tqdm = tracked_tqdm + + # Patch in sys.modules (same as HFProgressTracker) + tqdm_attr_names = ['tqdm', 'base_tqdm', 'old_tqdm'] + patched_count = 0 + + for module_name in list(sys.modules.keys()): + if "huggingface" in module_name or module_name.startswith("tqdm"): + try: + module = sys.modules[module_name] + for attr_name in tqdm_attr_names: + if hasattr(module, attr_name): + attr = getattr(module, attr_name) + is_tqdm_class = ( + attr is self._original_tqdm_class + or (self._original_tqdm_auto and attr is self._original_tqdm_auto) + or ( + hasattr(attr, "__name__") + and attr.__name__ == "tqdm" + and hasattr(attr, "update") + ) + ) + if is_tqdm_class: + key = f"{module_name}.{attr_name}" + self._patched_modules[key] = (module, attr_name, attr) + setattr(module, attr_name, tracked_tqdm) + patched_count += 1 + except (AttributeError, TypeError): + pass + + # Monkey-patch HF's tqdm.update (same as HFProgressTracker) + try: + from huggingface_hub.utils import tqdm as hf_tqdm_module + if hasattr(hf_tqdm_module, 'tqdm'): + hf_tqdm_class = hf_tqdm_module.tqdm + self._hf_tqdm_original_update = hf_tqdm_class.update + spy = self + + def patched_update(tqdm_self, n=1): + result = spy._hf_tqdm_original_update(tqdm_self, n) + desc = getattr(tqdm_self, 'desc', '') or '' + current = getattr(tqdm_self, 'n', 0) + total = getattr(tqdm_self, 'total', 0) or 0 + + spy._log( + "HF_UPDATE", + desc=desc, + current=current, + total=total, + pct=f"{100 * current / total:.1f}%" if total else "?", + ) + return result + + hf_tqdm_class.update = patched_update + patched_count += 1 + except (ImportError, AttributeError): + pass + + print(f"\n=== Patched {patched_count} tqdm references ===\n", flush=True) + + try: + yield + finally: + # Restore everything + import tqdm as tqdm_module + tqdm_module.tqdm = self._original_tqdm_class + if self._original_tqdm_auto: + tqdm_module.auto.tqdm = self._original_tqdm_auto + for key, (module, attr_name, original) in self._patched_modules.items(): + try: + setattr(module, attr_name, original) + except (AttributeError, TypeError): + pass + if self._hf_tqdm_original_update: + try: + from huggingface_hub.utils import tqdm as hf_tqdm_module + if hasattr(hf_tqdm_module, 'tqdm'): + hf_tqdm_module.tqdm.update = self._hf_tqdm_original_update + except (ImportError, AttributeError): + pass + + def summary(self): + print("\n" + "=" * 70) + print("SUMMARY") + print("=" * 70) + + inits = [e for e in self.events if e["type"] == "INIT"] + updates = [e for e in self.events if e["type"] in ("UPDATE", "HF_UPDATE")] + + print(f"\ntqdm bars created: {len(inits)}") + for e in inits: + print(f" - {e.get('filename', '?'):40s} total={e.get('total', '?')}") + + print(f"\nTotal update calls: {len(updates)}") + + # Group updates by filename + by_file = {} + for e in updates: + fn = e.get("filename") or e.get("desc", "unknown") + if fn not in by_file: + by_file[fn] = [] + by_file[fn].append(e) + + for fn, evts in by_file.items(): + max_current = max(e.get("current", 0) for e in evts) + max_total = max(e.get("total", 0) for e in evts) + print(f"\n {fn}:") + print(f" updates: {len(evts)}") + print(f" max current: {max_current:,}") + print(f" max total: {max_total:,}") + if max_total > 0 and max_current > 0: + print(f" final pct: {100 * max_current / max_total:.1f}%") + else: + print(f" final pct: NO PROGRESS REPORTED") + + +# ─── Delete cache ───────────────────────────────────────────────────────────── + +def delete_cache(repo_id: str): + from huggingface_hub import constants as hf_constants + cache_dir = Path(hf_constants.HF_HUB_CACHE) + repo_cache = cache_dir / ("models--" + repo_id.replace("/", "--")) + if repo_cache.exists(): + print(f"Deleting cache: {repo_cache}") + shutil.rmtree(repo_cache) + print("Deleted.") + else: + print(f"No cache found at {repo_cache}") + + +# ─── Download functions ─────────────────────────────────────────────────────── + +def download_qwen(spy: ProgressSpy): + """Mirrors how pytorch_backend.py downloads Qwen.""" + from transformers import AutoModel + repo_id = MODELS["qwen"]["repo_id"] + + print(f"Downloading {repo_id} via AutoModel.from_pretrained...") + with spy.patch(): + # This is what Qwen3TTSModel.from_pretrained does under the hood + from huggingface_hub import snapshot_download + snapshot_download(repo_id) + + +def download_luxtts(spy: ProgressSpy): + """Mirrors how luxtts_backend.py downloads LuxTTS.""" + from huggingface_hub import snapshot_download + repo_id = MODELS["luxtts"]["repo_id"] + + print(f"Downloading {repo_id} via snapshot_download...") + with spy.patch(): + snapshot_download(repo_id) + + +def download_chatterbox(spy: ProgressSpy): + """Mirrors how chatterbox_backend.py downloads Chatterbox.""" + from huggingface_hub import snapshot_download + cfg = MODELS["chatterbox"] + + print(f"Downloading {cfg['repo_id']} via snapshot_download with allow_patterns...") + with spy.patch(): + snapshot_download( + repo_id=cfg["repo_id"], + repo_type="model", + revision="main", + allow_patterns=cfg["allow_patterns"], + token=os.getenv("HF_TOKEN"), + ) + + +# ─── Main ───────────────────────────────────────────────────────────────────── + +def main(): + if len(sys.argv) < 2 or sys.argv[1] not in MODELS: + print(f"Usage: {sys.argv[0]} <{'|'.join(MODELS.keys())}> [--delete]") + sys.exit(1) + + model_key = sys.argv[1] + should_delete = "--delete" in sys.argv + cfg = MODELS[model_key] + + print(f"\n{'=' * 70}") + print(f"Testing download progress for: {cfg['description']}") + print(f"Repo: {cfg['repo_id']}") + print(f"Method: {cfg['method']}") + print(f"{'=' * 70}\n") + + if should_delete: + delete_cache(cfg["repo_id"]) + print() + + spy = ProgressSpy() + + dispatch = { + "qwen": download_qwen, + "luxtts": download_luxtts, + "chatterbox": download_chatterbox, + } + + try: + dispatch[model_key](spy) + except Exception as e: + print(f"\n!!! Download failed: {e}") + + spy.summary() + + +if __name__ == "__main__": + main() diff --git a/tauri/src-tauri/gen/Assets.car b/tauri/src-tauri/gen/Assets.car index de0e9a08..a5f68f95 100644 Binary files a/tauri/src-tauri/gen/Assets.car and b/tauri/src-tauri/gen/Assets.car differ diff --git a/tauri/src-tauri/gen/voicebox.icns b/tauri/src-tauri/gen/voicebox.icns index 59661d99..e4492f52 100644 Binary files a/tauri/src-tauri/gen/voicebox.icns and b/tauri/src-tauri/gen/voicebox.icns differ