--- title: "Model Management" description: "How model downloading, loading, and status tracking works in Voicebox" --- ## Overview Voicebox manages two types of models: **TTS Models:** Qwen3-TTS for voice cloning (0.6B and 1.7B variants). **ASR Models:** Whisper for transcription (tiny through large). Models are downloaded from HuggingFace Hub on first use and cached locally. ## Available Models ### TTS Models | Model | HuggingFace ID | Size | VRAM | |-------|----------------|------|------| | 0.6B | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ~1.2GB | ~2GB | | 1.7B | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ~3.4GB | ~6GB | ### Whisper Models | Model | HuggingFace ID | Size | VRAM | |-------|----------------|------|------| | tiny | `openai/whisper-tiny` | ~150MB | ~1GB | | base | `openai/whisper-base` | ~300MB | ~1GB | | small | `openai/whisper-small` | ~500MB | ~2GB | | medium | `openai/whisper-medium` | ~1.5GB | ~5GB | | large | `openai/whisper-large` | ~3GB | ~10GB | ## Model Storage Models are cached in the HuggingFace cache directory: ``` ~/.cache/huggingface/hub/ ├── models--Qwen--Qwen3-TTS-12Hz-1.7B-Base/ ├── models--Qwen--Qwen3-TTS-12Hz-0.6B-Base/ ├── models--openai--whisper-base/ └── ... ``` ## Progress Tracking ### Progress Manager Tracks download progress across all models: ```python class ProgressManager: def __init__(self): self._progress = {} # model_name -> progress_info def update_progress( self, model_name: str, current: int, total: int, filename: str, status: str, ): self._progress[model_name] = { "current": current, "total": total, "filename": filename, "status": status, # downloading, complete, error "updated_at": datetime.utcnow(), } def get_progress(self, model_name: str) -> Optional[dict]: return self._progress.get(model_name) ``` ### HuggingFace Progress Callback Hooks into HuggingFace's download system: ```python class HFProgressTracker: def __init__(self, callback): self.callback = callback @contextmanager def patch_download(self): """Context manager to intercept HF downloads.""" original_download = hf_hub_download def patched_download(*args, **kwargs): # Intercept progress result = original_download(*args, **kwargs) self.callback(progress_info) return result # Apply patch with patch('huggingface_hub.hf_hub_download', patched_download): yield ``` ### Server-Sent Events (SSE) Progress is streamed to the frontend: ```python @app.get("/models/progress/{model_name}") async def get_model_progress(model_name: str): async def event_generator(): while True: progress = progress_manager.get_progress(model_name) if progress: yield f"data: {json.dumps(progress)}\n\n" if progress and progress["status"] in ["complete", "error"]: break await asyncio.sleep(0.5) return StreamingResponse( event_generator(), media_type="text/event-stream" ) ``` ## Task Manager Tracks active downloads and generations: ```python class TaskManager: def __init__(self): self._active_downloads = {} self._active_generations = {} def start_download(self, model_name: str): self._active_downloads[model_name] = { "status": "downloading", "started_at": datetime.utcnow(), } def complete_download(self, model_name: str): if model_name in self._active_downloads: del self._active_downloads[model_name] def get_active_tasks(self) -> dict: return { "downloads": list(self._active_downloads.values()), "generations": list(self._active_generations.values()), } ``` ## Model Status Check which models are downloaded and loaded: ```python @app.get("/models/status") async def get_model_status() -> ModelStatusListResponse: models = [] # Check TTS models for size, hf_id in [("1.7B", "Qwen/Qwen3-TTS-12Hz-1.7B-Base"), ...]: downloaded = is_model_downloaded(hf_id) loaded = tts_model._current_model_size == size models.append(ModelStatus( model_name=f"qwen-tts-{size}", display_name=f"Qwen3-TTS {size}", downloaded=downloaded, size_mb=get_model_size_mb(hf_id), loaded=loaded, )) # Check Whisper models for size in ["tiny", "base", "small", "medium", "large"]: hf_id = f"openai/whisper-{size}" downloaded = is_model_downloaded(hf_id) models.append(ModelStatus( model_name=f"whisper-{size}", display_name=f"Whisper {size}", downloaded=downloaded, size_mb=get_model_size_mb(hf_id), loaded=False, # Whisper is loaded on-demand )) return ModelStatusListResponse(models=models) ``` ## Manual Model Operations ### Load Model ```python @app.post("/models/load") async def load_model(model_size: str = "1.7B"): tts_model = get_tts_model() await tts_model.load_model_async(model_size) return {"status": "loaded", "model_size": model_size} ``` ### Unload Model ```python @app.post("/models/unload") async def unload_model(): tts_model = get_tts_model() tts_model.unload_model() return {"status": "unloaded"} ``` ### Trigger Download ```python @app.post("/models/download") async def trigger_model_download(request: ModelDownloadRequest): # This triggers the download in background # Progress is tracked via /models/progress/{model_name} if request.model_name.startswith("qwen-tts"): size = request.model_name.split("-")[-1] asyncio.create_task(download_tts_model(size)) elif request.model_name.startswith("whisper"): size = request.model_name.split("-")[-1] asyncio.create_task(download_whisper_model(size)) return {"status": "downloading"} ``` ### Delete Model ```python @app.delete("/models/{model_name}") async def delete_model(model_name: str): # Find and delete from HuggingFace cache cache_dir = Path.home() / ".cache" / "huggingface" / "hub" model_dirs = list(cache_dir.glob(f"models--*--{model_name}*")) for model_dir in model_dirs: shutil.rmtree(model_dir) return {"status": "deleted"} ``` ## API Endpoints | Method | Endpoint | Description | |--------|----------|-------------| | GET | `/models/status` | Get status of all models | | POST | `/models/load` | Load TTS model | | POST | `/models/unload` | Unload TTS model | | POST | `/models/download` | Trigger model download | | GET | `/models/progress/{name}` | Stream download progress (SSE) | | DELETE | `/models/{name}` | Delete downloaded model | | GET | `/tasks/active` | Get active downloads/generations | ## Response Schemas ### ModelStatus ```json { "model_name": "qwen-tts-1.7B", "display_name": "Qwen3-TTS 1.7B", "downloaded": true, "size_mb": 3400, "loaded": true } ``` ### ActiveTasksResponse ```json { "downloads": [ { "model_name": "whisper-medium", "status": "downloading", "started_at": "2024-01-15T10:30:00Z" } ], "generations": [ { "task_id": "uuid", "profile_id": "uuid", "text_preview": "Hello world...", "started_at": "2024-01-15T10:30:00Z" } ] } ``` ## Frontend Integration ### Progress Display ```typescript // Subscribe to download progress via SSE const eventSource = new EventSource(`/models/progress/${modelName}`); eventSource.onmessage = (event) => { const progress = JSON.parse(event.data); updateProgressBar(progress.current / progress.total); if (progress.status === 'complete') { eventSource.close(); } }; ``` ### Model Status UI ```typescript // Fetch model status const { data: models } = useQuery({ queryKey: ['models', 'status'], queryFn: () => api.getModelStatus(), }); // Display download/load buttons based on status models.map(model => ( triggerDownload(model.model_name)} onLoad={() => loadModel(model.model_name)} /> )); ``` ## Error Handling | Error | Cause | Solution | |-------|-------|----------| | Download failed | Network issue | Retry download | | OOM on load | Model too large | Use smaller model | | Model not found | Cache corrupted | Re-download | | Slow download | HF rate limit | Wait and retry |