diff --git a/app/src/App.tsx b/app/src/App.tsx index 69bdd277..ed912a55 100644 --- a/app/src/App.tsx +++ b/app/src/App.tsx @@ -1,63 +1,51 @@ -import { History, Mic, Settings, Sparkles } from 'lucide-react'; +import { useState } from 'react'; import { GenerationForm } from '@/components/Generation/GenerationForm'; import { HistoryTable } from '@/components/History/HistoryTable'; import { ConnectionForm } from '@/components/ServerSettings/ConnectionForm'; import { ServerStatus } from '@/components/ServerSettings/ServerStatus'; -import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs'; +import { ModelManagement } from '@/components/ServerSettings/ModelManagement'; import { Toaster } from '@/components/ui/toaster'; import { ProfileList } from '@/components/VoiceProfiles/ProfileList'; +import { Sidebar } from '@/components/Sidebar'; function App() { + const [activeTab, setActiveTab] = useState('profiles'); + return ( -
-
-
-

voicebox

-

- Production-quality Qwen3-TTS voice cloning and generation -

-
- - - - - - Profiles - - - - Generate - - - - History - - - - Settings - - - - - - - - - - - - - - - - -
- - +
+ + +
+
+ {activeTab === 'profiles' && ( +
+
- - -
+ )} + + {activeTab === 'generate' && ( +
+ +
+ )} + + {activeTab === 'history' && ( +
+ +
+ )} + + {activeTab === 'settings' && ( +
+
+ + +
+ +
+ )} +
+
diff --git a/app/src/components/ServerSettings/ModelManagement.tsx b/app/src/components/ServerSettings/ModelManagement.tsx new file mode 100644 index 00000000..1c9815bd --- /dev/null +++ b/app/src/components/ServerSettings/ModelManagement.tsx @@ -0,0 +1,179 @@ +import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query'; +import { apiClient } from '@/lib/api/client'; +import { Card, CardContent, CardHeader, CardTitle, CardDescription } from '@/components/ui/card'; +import { Button } from '@/components/ui/button'; +import { Badge } from '@/components/ui/badge'; +import { Loader2, Download, CheckCircle2 } from 'lucide-react'; +import { ModelProgress } from './ModelProgress'; +import { useToast } from '@/components/ui/use-toast'; + +export function ModelManagement() { + const { toast } = useToast(); + const queryClient = useQueryClient(); + + const { data: modelStatus, isLoading } = useQuery({ + queryKey: ['modelStatus'], + queryFn: () => apiClient.getModelStatus(), + refetchInterval: 5000, // Refresh every 5 seconds + }); + + const downloadMutation = useMutation({ + mutationFn: (modelName: string) => apiClient.triggerModelDownload(modelName), + onSuccess: (_, modelName) => { + toast({ + title: 'Download started', + description: `Downloading ${modelName}...`, + }); + // Refetch status after a delay to see progress + setTimeout(() => { + queryClient.invalidateQueries({ queryKey: ['modelStatus'] }); + }, 1000); + }, + onError: (error: Error) => { + toast({ + title: 'Download failed', + description: error.message, + variant: 'destructive', + }); + }, + }); + + const formatSize = (sizeMb?: number): string => { + if (!sizeMb) return 'Unknown'; + if (sizeMb < 1024) return `${sizeMb.toFixed(1)} MB`; + return `${(sizeMb / 1024).toFixed(2)} GB`; + }; + + return ( + + + Model Management + + 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) => ( + downloadMutation.mutate(model.model_name)} + isDownloading={downloadMutation.isPending} + formatSize={formatSize} + /> + ))} +
+
+ + {/* Whisper Models */} +
+

Transcription Models

+
+ {modelStatus.models + .filter((m) => m.model_name.startsWith('whisper')) + .map((model) => ( + downloadMutation.mutate(model.model_name)} + isDownloading={downloadMutation.isPending} + formatSize={formatSize} + /> + ))} +
+
+ + {/* Progress indicators */} +
+

Download Progress

+
+ {modelStatus.models.map((model) => ( + + ))} +
+
+
+ ) : null} + + + ); +} + +interface ModelItemProps { + model: { + model_name: string; + display_name: string; + downloaded: boolean; + size_mb?: number; + loaded: boolean; + }; + onDownload: () => void; + isDownloading: boolean; + formatSize: (sizeMb?: number) => string; +} + +function ModelItem({ model, onDownload, isDownloading, formatSize }: ModelItemProps) { + return ( +
+
+
+ {model.display_name} + {model.loaded && ( + Loaded + )} + {model.downloaded && !model.loaded && ( + Downloaded + )} +
+ {model.downloaded && model.size_mb && ( +
+ Size: {formatSize(model.size_mb)} +
+ )} +
+
+ {model.downloaded ? ( +
+ + Ready +
+ ) : ( + + )} +
+
+ ); +} diff --git a/app/src/components/ServerSettings/ModelProgress.tsx b/app/src/components/ServerSettings/ModelProgress.tsx new file mode 100644 index 00000000..79ae1f0e --- /dev/null +++ b/app/src/components/ServerSettings/ModelProgress.tsx @@ -0,0 +1,121 @@ +import { useEffect, useState } from 'react'; +import { Progress } from '@/components/ui/progress'; +import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card'; +import { useServerStore } from '@/stores/serverStore'; +import type { ModelProgress as ModelProgressType } from '@/lib/api/types'; +import { Loader2, CheckCircle2, XCircle } from 'lucide-react'; + +interface ModelProgressProps { + modelName: string; + displayName: string; +} + +export function ModelProgress({ modelName, displayName }: ModelProgressProps) { + const [progress, setProgress] = useState(null); + const [isSubscribed, setIsSubscribed] = useState(false); + const serverUrl = useServerStore((state) => state.serverUrl); + + useEffect(() => { + if (!serverUrl || isSubscribed) return; + + // Subscribe to progress updates via Server-Sent Events + const eventSource = new EventSource(`${serverUrl}/models/progress/${modelName}`); + + eventSource.onmessage = (event) => { + try { + const data = JSON.parse(event.data) as ModelProgressType; + setProgress(data); + + // Close connection if complete or error + if (data.status === 'complete' || data.status === 'error') { + eventSource.close(); + setIsSubscribed(false); + } + } catch (error) { + console.error('Error parsing progress event:', error); + } + }; + + eventSource.onerror = (error) => { + console.error('SSE error:', error); + eventSource.close(); + setIsSubscribed(false); + }; + + setIsSubscribed(true); + + return () => { + eventSource.close(); + setIsSubscribed(false); + }; + }, [serverUrl, modelName, isSubscribed]); + + // Don't render if no progress or if complete/error and some time has passed + if (!progress || (progress.status === 'complete' && Date.now() - new Date(progress.timestamp).getTime() > 5000)) { + return null; + } + + const 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 / Math.pow(k, i)).toFixed(1)} ${sizes[i]}`; + }; + + const getStatusIcon = () => { + switch (progress.status) { + case 'complete': + return ; + case 'error': + return ; + case 'downloading': + case 'extracting': + return ; + default: + return null; + } + }; + + const getStatusText = () => { + switch (progress.status) { + case 'complete': + return 'Download complete'; + case 'error': + return `Error: ${progress.error || 'Unknown error'}`; + case 'downloading': + return progress.filename ? `Downloading ${progress.filename}...` : 'Downloading...'; + case 'extracting': + return 'Extracting...'; + default: + return 'Processing...'; + } + }; + + return ( + + + + {getStatusIcon()} + {displayName} + + + +
+
+ {getStatusText()} + {progress.total > 0 && ( + + {formatBytes(progress.current)} / {formatBytes(progress.total)} ( + {progress.progress.toFixed(1)}%) + + )} +
+ {progress.total > 0 && ( + + )} +
+ + + ); +} diff --git a/app/src/components/ServerSettings/ServerStatus.tsx b/app/src/components/ServerSettings/ServerStatus.tsx index 7b0ad9ee..e37fa5b3 100644 --- a/app/src/components/ServerSettings/ServerStatus.tsx +++ b/app/src/components/ServerSettings/ServerStatus.tsx @@ -3,6 +3,7 @@ import { Badge } from '@/components/ui/badge'; import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card'; import { useServerHealth } from '@/lib/hooks/useServer'; import { useServerStore } from '@/stores/serverStore'; +import { ModelProgress } from './ModelProgress'; export function ServerStatus() { const { data: health, isLoading, error } = useServerHealth(); @@ -19,6 +20,16 @@ export function ServerStatus() {
{serverUrl}
+ {/* Model download progress */} +
+ + + + + + +
+ {isLoading ? (
@@ -35,10 +46,19 @@ export function ServerStatus() { Connected
-
+
- Model: {health.model_loaded ? 'Loaded' : 'Not Loaded'} + Model: {health.model_loaded + ? `Loaded${health.model_size ? ` (${health.model_size})` : ''}` + : health.model_downloaded === false + ? 'Not Downloaded' + : 'Not Loaded'} + {health.model_downloaded === true && !health.model_loaded && ( + + Model Cached (will load on first use) + + )} GPU: {health.gpu_available ? 'Available' : 'Not Available'} diff --git a/app/src/components/Sidebar.tsx b/app/src/components/Sidebar.tsx new file mode 100644 index 00000000..400976bd --- /dev/null +++ b/app/src/components/Sidebar.tsx @@ -0,0 +1,47 @@ +import { History, Mic, Settings, Sparkles } from 'lucide-react'; +import { cn } from '@/lib/utils/cn'; + +interface SidebarProps { + activeTab: string; + onTabChange: (tab: string) => void; +} + +const tabs = [ + { id: 'profiles', icon: Mic, label: 'Profiles' }, + { id: 'generate', icon: Sparkles, label: 'Generate' }, + { id: 'history', icon: History, label: 'History' }, + { id: 'settings', icon: Settings, label: 'Settings' }, +]; + +export function Sidebar({ activeTab, onTabChange }: SidebarProps) { + return ( +
+ {/* Navigation Buttons */} +
+ {tabs.map((tab) => { + const Icon = tab.icon; + const isActive = activeTab === tab.id; + + return ( + + ); + })} +
+
+ ); +} diff --git a/app/src/index.css b/app/src/index.css index aaf31c1f..a587db92 100644 --- a/app/src/index.css +++ b/app/src/index.css @@ -32,6 +32,7 @@ --color-border: hsl(var(--border)); --color-input: hsl(var(--input)); --color-ring: hsl(var(--ring)); + --color-sidebar: hsl(var(--sidebar)); --color-chart-1: hsl(var(--chart-1)); --color-chart-2: hsl(var(--chart-2)); @@ -60,6 +61,7 @@ --border: 214.3 31.8% 91.4%; --input: 214.3 31.8% 91.4%; --ring: 222.2 84% 4.9%; + --sidebar: 0 0% 98%; --radius: 0.5rem; --chart-1: 12 76% 61%; --chart-2: 173 58% 39%; @@ -69,25 +71,26 @@ } .dark { - --background: 222.2 84% 4.9%; - --foreground: 210 40% 98%; - --card: 222.2 84% 4.9%; - --card-foreground: 210 40% 98%; - --popover: 222.2 84% 4.9%; - --popover-foreground: 210 40% 98%; - --primary: 210 40% 98%; - --primary-foreground: 222.2 47.4% 11.2%; - --secondary: 217.2 32.6% 17.5%; - --secondary-foreground: 210 40% 98%; - --muted: 217.2 32.6% 17.5%; - --muted-foreground: 215 20.2% 65.1%; - --accent: 217.2 32.6% 17.5%; - --accent-foreground: 210 40% 98%; - --destructive: 0 62.8% 30.6%; - --destructive-foreground: 210 40% 98%; - --border: 217.2 32.6% 17.5%; - --input: 217.2 32.6% 17.5%; - --ring: 212.7 26.8% 83.9%; + --background: 0 0% 8%; + --foreground: 0 0% 95%; + --card: 0 0% 10%; + --card-foreground: 0 0% 95%; + --popover: 0 0% 10%; + --popover-foreground: 0 0% 95%; + --primary: 0 0% 20%; + --primary-foreground: 0 0% 95%; + --secondary: 0 0% 15%; + --secondary-foreground: 0 0% 95%; + --muted: 0 0% 15%; + --muted-foreground: 0 0% 60%; + --accent: 0 0% 15%; + --accent-foreground: 0 0% 95%; + --destructive: 0 62.8% 50%; + --destructive-foreground: 0 0% 95%; + --border: 0 0% 15%; + --input: 0 0% 15%; + --ring: 0 0% 40%; + --sidebar: 0 0% 6%; --chart-1: 220 70% 50%; --chart-2: 160 60% 45%; --chart-3: 30 80% 55%; @@ -103,3 +106,11 @@ @apply bg-background text-foreground; } } + +@layer utilities { + .writing-vertical { + writing-mode: vertical-rl; + text-orientation: mixed; + letter-spacing: 0.1em; + } +} diff --git a/backend/main.py b/backend/main.py index 1a624c35..06bd366a 100644 --- a/backend/main.py +++ b/backend/main.py @@ -19,6 +19,7 @@ import uuid from . import database, models, profiles, history, tts, transcribe from .database import get_db, init_db, Generation as DBGeneration, VoiceProfile as DBVoiceProfile +from .utils.progress import get_progress_manager # Initialize database init_db() @@ -52,6 +53,10 @@ async def root(): @app.get("/health", response_model=models.HealthResponse) async def health(): """Health check endpoint.""" + from huggingface_hub import hf_hub_download + from pathlib import Path + import os + tts_model = tts.get_tts_model() gpu_available = torch.cuda.is_available() @@ -59,9 +64,44 @@ async def health(): if gpu_available: vram_used = torch.cuda.memory_allocated() / 1024 / 1024 # MB + # Check if model is loaded + model_loaded = tts_model.is_loaded() + model_size = tts_model.model_size if model_loaded else None + + # Check if default model is downloaded (cached) + model_downloaded = None + try: + # Check if the default model (1.7B) is cached + default_model_id = "Qwen/Qwen3-TTS-12Hz-1.7B-Base" + + # Method 1: Try scan_cache_dir if available + try: + from huggingface_hub import scan_cache_dir + cache_info = scan_cache_dir() + for repo in cache_info.repos: + if repo.repo_id == default_model_id: + model_downloaded = True + break + except (ImportError, Exception): + # Method 2: Check cache directory + cache_dir = os.path.expanduser("~/.cache/huggingface/hub") + repo_cache = Path(cache_dir) / "models--" + default_model_id.replace("/", "--") + if repo_cache.exists(): + has_model_files = ( + any(repo_cache.rglob("*.bin")) or + any(repo_cache.rglob("*.safetensors")) or + any(repo_cache.rglob("*.pt")) or + any(repo_cache.rglob("*.pth")) + ) + model_downloaded = has_model_files + except Exception: + pass + return models.HealthResponse( status="healthy", - model_loaded=tts_model.is_loaded(), + model_loaded=model_loaded, + model_downloaded=model_downloaded, + model_size=model_size, gpu_available=gpu_available, vram_used_mb=vram_used, ) @@ -395,6 +435,255 @@ async def unload_model(): raise HTTPException(status_code=500, detail=str(e)) +@app.get("/models/progress/{model_name}") +async def get_model_progress(model_name: str): + """Get model download progress via Server-Sent Events.""" + from fastapi.responses import StreamingResponse + + progress_manager = get_progress_manager() + + async def event_generator(): + """Generate SSE events for progress updates.""" + async for event in progress_manager.subscribe(model_name): + yield event + + return StreamingResponse( + event_generator(), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + + +@app.get("/models/status", response_model=models.ModelStatusListResponse) +async def get_model_status(): + """Get status of all available models.""" + from huggingface_hub import hf_hub_download + from pathlib import Path + import os + + # Try to import scan_cache_dir (might not be available in older versions) + try: + from huggingface_hub import scan_cache_dir + use_scan_cache = True + except ImportError: + use_scan_cache = False + + def check_tts_loaded(model_size: str): + """Check if TTS model is loaded with specific size.""" + try: + tts_model = tts.get_tts_model() + return tts_model.is_loaded() and tts_model.model_size == model_size + except Exception: + return False + + def check_whisper_loaded(model_size: str): + """Check if Whisper model is loaded with specific size.""" + try: + whisper_model = transcribe.get_whisper_model() + return whisper_model.is_loaded() and whisper_model.model_size == model_size + except Exception: + return False + + model_configs = [ + { + "model_name": "qwen-tts-1.7B", + "display_name": "Qwen TTS 1.7B", + "hf_repo_id": "Qwen/Qwen3-TTS-12Hz-1.7B-Base", + "model_size": "1.7B", + "check_loaded": lambda: check_tts_loaded("1.7B"), + }, + { + "model_name": "qwen-tts-0.6B", + "display_name": "Qwen TTS 0.6B", + "hf_repo_id": "Qwen/Qwen3-TTS-12Hz-0.6B-Base", + "model_size": "0.6B", + "check_loaded": lambda: check_tts_loaded("0.6B"), + }, + { + "model_name": "whisper-base", + "display_name": "Whisper Base", + "hf_repo_id": "openai/whisper-base", + "model_size": "base", + "check_loaded": lambda: check_whisper_loaded("base"), + }, + { + "model_name": "whisper-small", + "display_name": "Whisper Small", + "hf_repo_id": "openai/whisper-small", + "model_size": "small", + "check_loaded": lambda: check_whisper_loaded("small"), + }, + { + "model_name": "whisper-medium", + "display_name": "Whisper Medium", + "hf_repo_id": "openai/whisper-medium", + "model_size": "medium", + "check_loaded": lambda: check_whisper_loaded("medium"), + }, + { + "model_name": "whisper-large", + "display_name": "Whisper Large", + "hf_repo_id": "openai/whisper-large", + "model_size": "large", + "check_loaded": lambda: check_whisper_loaded("large"), + }, + ] + + # Get HuggingFace cache info (if available) + cache_info = None + if use_scan_cache: + try: + cache_info = scan_cache_dir() + except Exception: + # Function failed, continue without it + pass + + statuses = [] + + for config in model_configs: + try: + downloaded = False + size_mb = None + loaded = False + + # Method 1: Try using scan_cache_dir if available + if cache_info: + repo_id = config["hf_repo_id"] + for repo in cache_info.repos: + if repo.repo_id == repo_id: + downloaded = True + # Calculate size from cache info + try: + total_size = sum(revision.size_on_disk for revision in repo.revisions) + size_mb = total_size / (1024 * 1024) + except Exception: + pass + break + + # Method 2: Fallback to checking cache directory directly + if not downloaded: + try: + cache_dir = os.path.expanduser("~/.cache/huggingface/hub") + repo_cache = Path(cache_dir) / "models--" + config["hf_repo_id"].replace("/", "--") + + if repo_cache.exists(): + # Check for model files (bin, safetensors, or other common model files) + has_model_files = ( + any(repo_cache.rglob("*.bin")) or + any(repo_cache.rglob("*.safetensors")) or + any(repo_cache.rglob("*.pt")) or + any(repo_cache.rglob("*.pth")) or + any(repo_cache.rglob("model.safetensors.index.json")) or + any(repo_cache.rglob("pytorch_model.bin.index.json")) + ) + + if has_model_files: + downloaded = True + # Calculate size + try: + total_size = sum(f.stat().st_size for f in repo_cache.rglob("*") if f.is_file()) + size_mb = total_size / (1024 * 1024) + except Exception: + pass + except Exception: + pass + + # Method 3: Try to check if model can be loaded locally (last resort) + if not downloaded: + try: + # Try to download with local_files_only=True to check if cached + hf_hub_download( + repo_id=config["hf_repo_id"], + filename="config.json", # Try a common file + local_files_only=True, + ) + downloaded = True + except Exception: + # File not found locally, model not downloaded + pass + + # Check if loaded in memory + try: + loaded = config["check_loaded"]() + except Exception: + loaded = False + + statuses.append(models.ModelStatus( + model_name=config["model_name"], + display_name=config["display_name"], + downloaded=downloaded, + size_mb=size_mb, + loaded=loaded, + )) + except Exception as e: + # If check fails, try to at least check if loaded + try: + loaded = config["check_loaded"]() + except Exception: + loaded = False + + statuses.append(models.ModelStatus( + model_name=config["model_name"], + display_name=config["display_name"], + downloaded=False, # Assume not downloaded if check failed + size_mb=None, + loaded=loaded, + )) + + return models.ModelStatusListResponse(models=statuses) + + +@app.post("/models/download") +async def trigger_model_download(request: models.ModelDownloadRequest): + """Trigger download of a specific model.""" + import asyncio + + model_configs = { + "qwen-tts-1.7B": { + "model_size": "1.7B", + "load_func": lambda: tts.get_tts_model().load_model("1.7B"), + }, + "qwen-tts-0.6B": { + "model_size": "0.6B", + "load_func": lambda: tts.get_tts_model().load_model("0.6B"), + }, + "whisper-base": { + "model_size": "base", + "load_func": lambda: transcribe.get_whisper_model().load_model("base"), + }, + "whisper-small": { + "model_size": "small", + "load_func": lambda: transcribe.get_whisper_model().load_model("small"), + }, + "whisper-medium": { + "model_size": "medium", + "load_func": lambda: transcribe.get_whisper_model().load_model("medium"), + }, + "whisper-large": { + "model_size": "large", + "load_func": lambda: transcribe.get_whisper_model().load_model("large"), + }, + } + + if request.model_name not in model_configs: + raise HTTPException(status_code=400, detail=f"Unknown model: {request.model_name}") + + config = model_configs[request.model_name] + + try: + # Trigger download by loading the model (which will download if not cached) + # Run in background to avoid blocking + await asyncio.to_thread(config["load_func"]) + + return {"message": f"Model {request.model_name} download started"} + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + # ============================================ # STARTUP & SHUTDOWN # ============================================ diff --git a/backend/models.py b/backend/models.py index 7bd48ddb..47f7ffc7 100644 --- a/backend/models.py +++ b/backend/models.py @@ -111,5 +111,26 @@ class HealthResponse(BaseModel): """Response model for health check.""" status: str model_loaded: bool + model_downloaded: Optional[bool] = None # Whether model is cached/downloaded + model_size: Optional[str] = None # Current model size if loaded gpu_available: bool vram_used_mb: Optional[float] = None + + +class ModelStatus(BaseModel): + """Response model for model status.""" + model_name: str + display_name: str + downloaded: bool + size_mb: Optional[float] = None + loaded: bool = False + + +class ModelStatusListResponse(BaseModel): + """Response model for model status list.""" + models: List[ModelStatus] + + +class ModelDownloadRequest(BaseModel): + """Request model for triggering model download.""" + model_name: str diff --git a/backend/transcribe.py b/backend/transcribe.py index c40dec7c..afd1a5f9 100644 --- a/backend/transcribe.py +++ b/backend/transcribe.py @@ -6,6 +6,8 @@ from typing import Optional, List, Dict import torch import numpy as np from pathlib import Path +from .utils.progress import get_progress_manager +from .utils.hf_progress import HFProgressTracker, create_hf_progress_callback class WhisperModel: @@ -48,18 +50,33 @@ class WhisperModel: model_name = f"openai/whisper-{model_size}" + # Set up progress tracking + progress_manager = get_progress_manager() + progress_model_name = f"whisper-{model_size}" + print(f"Loading Whisper model {model_size} on {self.device}...") - self.processor = WhisperProcessor.from_pretrained(model_name) - self.model = WhisperForConditionalGeneration.from_pretrained(model_name) - self.model.to(self.device) + # Set up progress callback + progress_callback = create_hf_progress_callback(progress_model_name, progress_manager) + tracker = HFProgressTracker(progress_callback) + # Use progress tracker during download + with tracker.patch_download(): + self.processor = WhisperProcessor.from_pretrained(model_name) + self.model = WhisperForConditionalGeneration.from_pretrained(model_name) + + self.model.to(self.device) self.model_size = model_size + # Mark as complete + progress_manager.mark_complete(progress_model_name) + print(f"Whisper model {model_size} loaded successfully") except Exception as e: print(f"Error loading Whisper model: {e}") + progress_manager = get_progress_manager() + progress_manager.mark_error(f"whisper-{model_size}", str(e)) raise def unload_model(self): diff --git a/backend/tts.py b/backend/tts.py index ae428747..852946ae 100644 --- a/backend/tts.py +++ b/backend/tts.py @@ -11,6 +11,8 @@ from pathlib import Path from .utils.cache import get_cache_key, get_cached_voice_prompt, cache_voice_prompt from .utils.audio import normalize_audio +from .utils.progress import get_progress_manager +from .utils.hf_progress import HFProgressTracker, create_hf_progress_callback class TTSModel: @@ -99,14 +101,37 @@ class TTSModel: # Get model path (local or HuggingFace Hub ID) model_path = self._get_model_path(model_size) - print(f"Loading TTS model {model_size} on {self.device}...") + # Set up progress tracking + progress_manager = get_progress_manager() + model_name = f"qwen-tts-{model_size}" - # Load the model - from_pretrained handles both local paths and HF Hub IDs - self.model = Qwen3TTSModel.from_pretrained( - model_path, - device_map=self.device, - torch_dtype=torch.float32 if self.device == "cpu" else torch.bfloat16, - ) + # Check if model is being downloaded from HuggingFace Hub + if model_path.startswith("Qwen/"): + print(f"Loading TTS model {model_size} on {self.device}...") + + # Set up progress callback + progress_callback = create_hf_progress_callback(model_name, progress_manager) + tracker = HFProgressTracker(progress_callback) + + # Use progress tracker during download + with tracker.patch_download(): + # Load the model - downloads will happen automatically with progress tracking + self.model = Qwen3TTSModel.from_pretrained( + model_path, + device_map=self.device, + torch_dtype=torch.float32 if self.device == "cpu" else torch.bfloat16, + ) + + # Mark as complete + progress_manager.mark_complete(model_name) + else: + # Local model, no download needed + print(f"Loading TTS model {model_size} on {self.device}...") + self.model = Qwen3TTSModel.from_pretrained( + model_path, + device_map=self.device, + torch_dtype=torch.float32 if self.device == "cpu" else torch.bfloat16, + ) self._current_model_size = model_size self.model_size = model_size @@ -115,10 +140,14 @@ class TTSModel: except ImportError as e: print(f"Error: qwen_tts package not found. Install with: pip install git+https://github.com/QwenLM/Qwen3-TTS.git") + progress_manager = get_progress_manager() + progress_manager.mark_error(f"qwen-tts-{model_size}", str(e)) raise except Exception as e: print(f"Error loading TTS model: {e}") print(f"Tip: The model will be automatically downloaded from HuggingFace Hub on first use.") + progress_manager = get_progress_manager() + progress_manager.mark_error(f"qwen-tts-{model_size}", str(e)) raise def unload_model(self): diff --git a/backend/utils/hf_progress.py b/backend/utils/hf_progress.py new file mode 100644 index 00000000..657b6d21 --- /dev/null +++ b/backend/utils/hf_progress.py @@ -0,0 +1,93 @@ +""" +HuggingFace Hub download progress tracking. +""" + +from typing import Optional, Callable +from contextlib import contextmanager +import threading + + +class HFProgressTracker: + """Tracks HuggingFace Hub download progress by intercepting hf_hub_download.""" + + def __init__(self, progress_callback: Optional[Callable] = None): + self.progress_callback = progress_callback + self._original_hf_hub_download = None + self._lock = threading.Lock() + self._total_downloaded = 0 + self._total_size = 0 + + def _tracked_hf_hub_download(self, *args, **kwargs): + """Wrapper for hf_hub_download with progress tracking.""" + import huggingface_hub + + # Get original callback if present + original_resume_callback = kwargs.get("resume_download", None) + + def combined_callback(downloaded: int, total: int): + """Combined callback that tracks progress.""" + # Update totals + with self._lock: + # Estimate: assume each file contributes equally + # This is a simplification - in reality we'd track per-file + if total > 0: + self._total_size = max(self._total_size, total) + self._total_downloaded = downloaded + + # Call original callback if present + if original_resume_callback: + original_resume_callback(downloaded, total) + + # Call our progress callback + if self.progress_callback: + with self._lock: + self.progress_callback(self._total_downloaded, self._total_size) + + # Replace callback + kwargs["resume_download"] = combined_callback + + # Call original download + return self._original_hf_hub_download(*args, **kwargs) + + @contextmanager + def patch_download(self): + """Context manager to patch hf_hub_download for progress tracking.""" + try: + import huggingface_hub + self._original_hf_hub_download = huggingface_hub.hf_hub_download + + # Reset totals + with self._lock: + self._total_downloaded = 0 + self._total_size = 0 + + # Patch the function + huggingface_hub.hf_hub_download = self._tracked_hf_hub_download + + yield + except ImportError: + # If huggingface_hub not available, just yield without patching + yield + finally: + # Restore original + if self._original_hf_hub_download: + try: + import huggingface_hub + huggingface_hub.hf_hub_download = self._original_hf_hub_download + except ImportError: + pass + + +def create_hf_progress_callback(model_name: str, progress_manager): + """Create a progress callback for HuggingFace downloads.""" + def callback(downloaded: int, total: int): + """Progress callback.""" + if total > 0: + progress_manager.update_progress( + model_name=model_name, + current=downloaded, + total=total, + filename="", + status="downloading", + ) + return callback diff --git a/backend/utils/progress.py b/backend/utils/progress.py new file mode 100644 index 00000000..68360c16 --- /dev/null +++ b/backend/utils/progress.py @@ -0,0 +1,164 @@ +""" +Progress tracking for model downloads using Server-Sent Events. +""" + +from typing import Optional, Callable, Dict +from fastapi.responses import StreamingResponse +import asyncio +import json +from datetime import datetime + + +class ProgressManager: + """Manages download progress for multiple models.""" + + def __init__(self): + self._progress: Dict[str, Dict] = {} + self._listeners: Dict[str, list] = {} + + def update_progress( + self, + model_name: str, + current: int, + total: int, + filename: Optional[str] = None, + status: str = "downloading", + ): + """ + Update progress for a model download. + + Args: + model_name: Name of the model (e.g., "qwen-tts-1.7B", "whisper-base") + current: Current bytes downloaded + total: Total bytes to download + filename: Current file being downloaded + status: Status string (downloading, extracting, complete, error) + """ + progress_pct = (current / total * 100) if total > 0 else 0 + + self._progress[model_name] = { + "model_name": model_name, + "current": current, + "total": total, + "progress": progress_pct, + "filename": filename, + "status": status, + "timestamp": datetime.now().isoformat(), + } + + # Notify all listeners + if model_name in self._listeners: + for queue in self._listeners[model_name]: + try: + queue.put_nowait(self._progress[model_name].copy()) + except asyncio.QueueFull: + pass + + def get_progress(self, model_name: str) -> Optional[Dict]: + """Get current progress for a model.""" + return self._progress.get(model_name) + + def create_progress_callback(self, model_name: str, filename: Optional[str] = None): + """ + Create a progress callback function for HuggingFace downloads. + + Args: + model_name: Name of the model + filename: Optional filename filter + + Returns: + Callback function + """ + def callback(progress: Dict): + """HuggingFace Hub progress callback.""" + if "total" in progress and "current" in progress: + current = progress.get("current", 0) + total = progress.get("total", 0) + file_name = progress.get("filename", filename) + + self.update_progress( + model_name=model_name, + current=current, + total=total, + filename=file_name, + status="downloading", + ) + + return callback + + async def subscribe(self, model_name: str): + """ + Subscribe to progress updates for a model. + + Yields progress updates as Server-Sent Events. + """ + queue = asyncio.Queue(maxsize=10) + + # Add to listeners + if model_name not in self._listeners: + self._listeners[model_name] = [] + self._listeners[model_name].append(queue) + + try: + # Send initial progress if available + if model_name in self._progress: + yield f"data: {json.dumps(self._progress[model_name])}\n\n" + + # Stream updates + while True: + try: + # Wait for update with timeout + progress = await asyncio.wait_for(queue.get(), timeout=1.0) + yield f"data: {json.dumps(progress)}\n\n" + + # Stop if complete or error + if progress.get("status") in ("complete", "error"): + break + except asyncio.TimeoutError: + # Send heartbeat + yield ": heartbeat\n\n" + continue + finally: + # Remove from listeners + if model_name in self._listeners: + self._listeners[model_name].remove(queue) + if not self._listeners[model_name]: + del self._listeners[model_name] + + def mark_complete(self, model_name: str): + """Mark a model download as complete.""" + if model_name in self._progress: + self._progress[model_name]["status"] = "complete" + self._progress[model_name]["progress"] = 100.0 + # Notify listeners + if model_name in self._listeners: + for queue in self._listeners[model_name]: + try: + queue.put_nowait(self._progress[model_name].copy()) + except asyncio.QueueFull: + pass + + def mark_error(self, model_name: str, error: str): + """Mark a model download as failed.""" + if model_name in self._progress: + self._progress[model_name]["status"] = "error" + self._progress[model_name]["error"] = error + # Notify listeners + if model_name in self._listeners: + for queue in self._listeners[model_name]: + try: + queue.put_nowait(self._progress[model_name].copy()) + except asyncio.QueueFull: + pass + + +# Global progress manager instance +_progress_manager: Optional[ProgressManager] = None + + +def get_progress_manager() -> ProgressManager: + """Get or create the global progress manager.""" + global _progress_manager + if _progress_manager is None: + _progress_manager = ProgressManager() + return _progress_manager