From 6429cb66737edd12a1afe9028c84e0adfc66a8b7 Mon Sep 17 00:00:00 2001 From: Jamie Pine Date: Sun, 25 Jan 2026 03:10:16 -0800 Subject: [PATCH] Implement sidebar navigation and model management features. Refactor App component to utilize a Sidebar for tab navigation, integrating ProfileList, GenerationForm, HistoryTable, and ServerStatus components. Introduce ModelManagement and ModelProgress components for handling AI model downloads and status updates. Enhance CSS for sidebar styling and add progress tracking functionality in the backend for model downloads. --- app/src/App.tsx | 88 +++--- .../ServerSettings/ModelManagement.tsx | 179 +++++++++++ .../ServerSettings/ModelProgress.tsx | 121 ++++++++ .../ServerSettings/ServerStatus.tsx | 24 +- app/src/components/Sidebar.tsx | 47 +++ app/src/index.css | 49 +-- backend/main.py | 291 +++++++++++++++++- backend/models.py | 21 ++ backend/transcribe.py | 23 +- backend/tts.py | 43 ++- backend/utils/hf_progress.py | 93 ++++++ backend/utils/progress.py | 164 ++++++++++ 12 files changed, 1061 insertions(+), 82 deletions(-) create mode 100644 app/src/components/ServerSettings/ModelManagement.tsx create mode 100644 app/src/components/ServerSettings/ModelProgress.tsx create mode 100644 app/src/components/Sidebar.tsx create mode 100644 backend/utils/hf_progress.py create mode 100644 backend/utils/progress.py 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