Spaces:
Sleeping
Sleeping
| """HTTP API behind the web UI. | |
| Encodes text queries with CLAP. The 2D projection and evaluation tables are | |
| computed once and cached to disk. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| from pathlib import Path | |
| import numpy as np | |
| from fastapi import FastAPI, HTTPException | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import FileResponse | |
| from fastapi.staticfiles import StaticFiles | |
| from pydantic import BaseModel, Field | |
| from . import onnx_encoder | |
| from .evaluate import embedding_health | |
| from .voices import ( | |
| FILTERABLE, | |
| VoiceCatalog, | |
| attribute_retrieval, | |
| facets, | |
| linear_probe, | |
| matching_indices, | |
| search_voices, | |
| ) | |
| # Each query targets a label the catalog carries, so precision@10 is measurable. | |
| PROBE_QUERIES = [ | |
| ("gender", "female", "a recording of a woman speaking"), | |
| ("gender", "male", "a recording of a man speaking"), | |
| ("age", "young", "a recording of a young person in their twenties speaking"), | |
| ("age", "old", "a recording of an elderly person speaking"), | |
| ] | |
| # About sound -- timbre, pace, energy -- because that is what CLAP encodes. | |
| # Accent and language go through the filters; prompting for them is near random. | |
| SUGGESTIONS = [ | |
| "a warm raspy middle aged male narrator for an audiobook", | |
| "a deep authoritative movie trailer voice", | |
| "a gentle soothing voice for a meditation app", | |
| "a bright energetic young woman for a commercial", | |
| "a calm unhurried storyteller", | |
| "an intense dramatic villain", | |
| ] | |
| class SearchRequest(BaseModel): | |
| query: str = Field(min_length=1, max_length=500) | |
| topN: int = Field(default=12, ge=1, le=50) | |
| # Applied before ranking. The library labels these exactly, and they are | |
| # hard constraints -- CLAP cannot hear accent or language at all, and hears | |
| # gender at 0.90-1.00, which is not a guarantee. | |
| filters: dict[str, str] = Field(default_factory=dict) | |
| def _encoder(model_id: str, models: str | Path | None): | |
| """Prefer the ONNX export; fall back to torch. | |
| Serving only ever encodes text, so the int8 text tower is enough. It is | |
| lossy -- see mss.onnx_encoder -- and absent unless `mss export` has run. | |
| """ | |
| if models and onnx_encoder.available(models): | |
| try: | |
| return onnx_encoder.OnnxTextEncoder(Path(models)), "onnx-int8" | |
| except ImportError: | |
| pass | |
| # Imported here, not at module scope: mss.clap pulls in torch, which is | |
| # ~200MB resident and the whole reason the ONNX path exists. | |
| from .clap import ClapEncoder | |
| return ClapEncoder(model_id=model_id), "torch" | |
| class State: | |
| """Loaded catalog plus the derived data the UI needs.""" | |
| def __init__(self, root: Path, models: str | Path | None = None) -> None: | |
| self.root = root | |
| self.catalog = VoiceCatalog.load(root) | |
| if not len(self.catalog) or self.catalog.vectors is None: | |
| raise RuntimeError(f"no voices at {root} -- run `mss fetch` first") | |
| self.encoder, self.backend = _encoder(self.catalog.model_id, models) | |
| self.by_id = {v.voice_id: i for i, v in enumerate(self.catalog.voices)} | |
| # -- cached derivations ------------------------------------------- | |
| def _cached(self, name: str, build): | |
| path = self.root / f"{name}.json" | |
| if path.exists(): | |
| return json.loads(path.read_text()) | |
| data = build() | |
| path.write_text(json.dumps(data)) | |
| return data | |
| def projection(self) -> dict: | |
| """2D t-SNE of the voice vectors, for the embedding-space view.""" | |
| def build(): | |
| from sklearn.manifold import TSNE | |
| xy = TSNE(n_components=2, perplexity=30, init="pca", | |
| random_state=0).fit_transform(self.catalog.vectors) | |
| # normalise to [0, 1] so the UI can scale to any viewport | |
| lo, hi = xy.min(axis=0), xy.max(axis=0) | |
| xy = (xy - lo) / np.maximum(hi - lo, 1e-9) | |
| return { | |
| "points": [ | |
| {"id": v.voice_id, "name": v.name, | |
| "x": round(float(p[0]), 4), "y": round(float(p[1]), 4), | |
| "gender": v.gender, "age": v.age, "use_case": v.use_case} | |
| for v, p in zip(self.catalog.voices, xy) | |
| ] | |
| } | |
| return self._cached("projection", build) | |
| def evaluation(self) -> dict: | |
| """Zero-shot, linear probe and retrieval precision, computed once.""" | |
| def build(): | |
| return { | |
| "probe": [linear_probe(self.catalog, a) | |
| for a in ("gender", "age", "use_case", "language")], | |
| "retrieval": [attribute_retrieval(self.catalog, self.encoder, a, v, q) | |
| for a, v, q in PROBE_QUERIES], | |
| } | |
| return self._cached("evaluation", build) | |
| def create_app(root: str | Path, models: str | Path | None = None) -> FastAPI: | |
| state = State(Path(root).expanduser(), models) | |
| app = FastAPI(title="Voice Library Search", version="1.0") | |
| # The Vite dev server runs on a different origin. | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["http://localhost:5173", "http://127.0.0.1:5173"], | |
| allow_methods=["GET", "POST"], | |
| allow_headers=["*"], | |
| ) | |
| def voice_payload(index: int, score: float | None = None, rank: int | None = None) -> dict: | |
| v = state.catalog.voices[index] | |
| return { | |
| "id": v.voice_id, "name": v.name, "gender": v.gender, "age": v.age, | |
| "accent": v.accent, "useCase": v.use_case, "language": v.language, | |
| "description": v.description, | |
| "audio": f"/api/voices/{v.voice_id}/audio", | |
| **({"score": round(score, 6)} if score is not None else {}), | |
| **({"rank": rank} if rank is not None else {}), | |
| } | |
| def catalog(): | |
| from collections import Counter | |
| cat = state.catalog | |
| health = embedding_health(cat.vectors) | |
| return { | |
| "count": len(cat), | |
| "model": cat.model_id, | |
| "backend": state.backend, | |
| "dim": int(cat.vectors.shape[1]), | |
| "health": health, | |
| "suggestions": SUGGESTIONS, | |
| "filterable": list(FILTERABLE), | |
| "facets": {k: v for k, v in facets(cat).items()}, | |
| "distributions": { | |
| attr: Counter(x for x in cat.labels(attr) if x).most_common(6) | |
| for attr in ("gender", "age", "language", "use_case") | |
| }, | |
| } | |
| def search(req: SearchRequest): | |
| query = req.query.strip() | |
| if not query: | |
| raise HTTPException(400, "empty query") | |
| unknown = set(req.filters) - set(FILTERABLE) | |
| if unknown: | |
| raise HTTPException(400, f"unknown filter(s): {', '.join(sorted(unknown))}") | |
| # computed once and passed down so the catalog is not walked twice | |
| keep = matching_indices(state.catalog, req.filters) | |
| hits = search_voices(state.catalog, state.encoder, query, | |
| top_n=req.topN, keep=keep) | |
| candidates = len(keep) | |
| results = [ | |
| voice_payload(state.by_id[h.metadata["voice_id"]], h.score, h.rank) | |
| for h in hits | |
| ] | |
| return { | |
| "query": query, | |
| "total": len(state.catalog), | |
| "candidates": candidates, | |
| "filters": {k: v for k, v in req.filters.items() if v}, | |
| "results": results, | |
| } | |
| def heatmap(req: SearchRequest): | |
| """Similarity of every voice to the query, for shading the projection. | |
| Unlike /api/search this keeps all scores rather than the top N. | |
| """ | |
| query = req.query.strip() | |
| if not query: | |
| raise HTTPException(400, "empty query") | |
| scores = state.catalog.vectors @ state.encoder.encode_text([query])[0] | |
| order = np.argsort(-scores)[: req.topN] | |
| return { | |
| "query": query, | |
| # aligned index-for-index with /api/projection points | |
| "scores": [round(float(s), 4) for s in scores], | |
| "min": round(float(scores.min()), 4), | |
| "max": round(float(scores.max()), 4), | |
| "top": [voice_payload(int(j), float(scores[j]), r) | |
| for r, j in enumerate(order, 1)], | |
| } | |
| def similar(voice_id: str, k: int = 6): | |
| if voice_id not in state.by_id: | |
| raise HTTPException(404, "unknown voice") | |
| i = state.by_id[voice_id] | |
| scores = state.catalog.vectors @ state.catalog.vectors[i] | |
| order = np.argsort(-scores)[1 : k + 1] # [0] is the voice itself | |
| return { | |
| "seed": voice_payload(i), | |
| "results": [voice_payload(int(j), float(scores[j]), r) | |
| for r, j in enumerate(order, 1)], | |
| } | |
| def audio(voice_id: str): | |
| if voice_id not in state.by_id: | |
| raise HTTPException(404, "unknown voice") | |
| v = state.catalog.voices[state.by_id[voice_id]] | |
| path = state.catalog.preview_path(v) | |
| if path is None: | |
| raise HTTPException(404, "preview missing") | |
| return FileResponse(path, media_type="audio/mpeg") | |
| def projection(): | |
| return state.projection() | |
| def evaluation(): | |
| return state.evaluation() | |
| def robustness(): | |
| """Benchmark results for the Evaluation tab, read from disk. | |
| Missing files return null so the UI can name the command to run. | |
| """ | |
| out = {} | |
| for key, filename in (("compound", "compound_degradation.json"),): | |
| path = state.root / filename | |
| out[key] = json.loads(path.read_text()) if path.exists() else None | |
| return out | |
| app.state.backend = state.backend | |
| # Serve the built UI when it exists, so one process runs everything. | |
| dist = Path(__file__).resolve().parent.parent.parent / "ui" / "dist" | |
| if dist.is_dir(): | |
| app.mount("/", StaticFiles(directory=dist, html=True), name="ui") | |
| return app | |
| def serve(root: str | Path, host: str = "127.0.0.1", port: int = 8787, | |
| models: str | Path | None = None) -> None: | |
| import uvicorn | |
| app = create_app(root, models) | |
| dist = Path(__file__).resolve().parent.parent.parent / "ui" / "dist" | |
| print(f"Voice Library Search API -> http://{host}:{port}") | |
| print(f"text encoder: {app.state.backend}") | |
| print("UI: " + (f"http://{host}:{port}" if dist.is_dir() | |
| else "run `npm run dev` in ui/ (http://localhost:5173)")) | |
| uvicorn.run(app, host=host, port=port, log_level="warning") | |