Upload folder using huggingface_hub

Dockerfile

@@ -1,17 +1,19 @@
-# Ox Translate - container for Hugging Face Spaces (or any Docker host)
+# SwitchLingo - container for Hugging Face Spaces (or any Docker host)
 FROM python:3.12-slim
-RUN apt-get update && apt-get install -y --no-install-recommends ffmpeg fonts-dejavu-core && rm -rf /var/lib/apt/lists/*
+RUN apt-get update && apt-get install -y --no-install-recommends ffmpeg fonts-dejavu-core fonts-noto-core fonts-noto-cjk && rm -rf /var/lib/apt/lists/*
 RUN useradd -m -u 1000 user
 USER user
 ENV HOME=/home/user PATH=/home/user/.local/bin:$PATH HF_HOME=/home/user/.cache/huggingface \
     OXT_DATA=/home/user/data PYTHONUNBUFFERED=1
 WORKDIR /home/user/app
 COPY --chown=user requirements.txt .
-RUN pip install --no-cache-dir --user -r requirements.txt
+RUN pip install --no-cache-dir --user torch --index-url https://download.pytorch.org/whl/cpu && \
+    pip install --no-cache-dir --user -r requirements.txt
 # Pre-download the speech + translation models so the first visitor doesn't wait
 RUN python -c "from faster_whisper import WhisperModel; WhisperModel('small', device='cpu', compute_type='int8')" && \
     python -c "from huggingface_hub import snapshot_download; snapshot_download('JustFrederik/nllb-200-distilled-600M-ct2-int8')" && \
-    python -c "import transformers; transformers.AutoTokenizer.from_pretrained('facebook/nllb-200-distilled-600M')"
+    python -c "import transformers; transformers.AutoTokenizer.from_pretrained('facebook/nllb-200-distilled-600M')" && \
+    python -c "from huggingface_hub import hf_hub_download as d; [d('myshell-ai/OpenVoiceV2', f) for f in ('converter/config.json', 'converter/checkpoint.pth')]"
 COPY --chown=user . .
 EXPOSE 7860
 CMD ["python", "-m", "uvicorn", "app.server:app", "--host", "0.0.0.0", "--port", "7860", "--proxy-headers", "--forwarded-allow-ips", "*"]

app/engines.py

@@ -175,11 +175,33 @@ def available_voices() -> list[str]:
     return out
 
 
+def eleven_clone(sample: Path) -> str:
+    """Instant voice clone from a speech sample; returns a temporary voice_id."""
+    import requests
+    with open(sample, "rb") as f:
+        r = requests.post("https://api.elevenlabs.io/v1/voices/add",
+                          headers={"xi-api-key": os.environ["ELEVENLABS_API_KEY"]},
+                          data={"name": f"switchlingo-tmp-{sample.parent.name}",
+                                "remove_background_noise": "true"},
+                          files={"files": (sample.name, f, "audio/mpeg")}, timeout=180)
+    r.raise_for_status()
+    return r.json()["voice_id"]
+
+
+def eleven_delete(voice_id: str) -> None:
+    import requests
+    try:
+        requests.delete(f"https://api.elevenlabs.io/v1/voices/{voice_id}",
+                        headers={"xi-api-key": os.environ["ELEVENLABS_API_KEY"]}, timeout=60)
+    except Exception:  # noqa: BLE001 - best effort cleanup
+        pass
+
+
 def speak(text: str, lang: str, out: Path, gender: str = "Female", engine: str = "edge",
-          rate: str = "+15%") -> None:
+          rate: str = "+15%", voice_id: str | None = None) -> None:
     """Write speech for `text` to `out` (mp3)."""
     if engine == "elevenlabs":
-        vid = os.environ.get(f"OXT_ELEVEN_VOICE_{gender.upper()}", ELEVEN_DEFAULT[gender])
+        vid = voice_id or os.environ.get(f"OXT_ELEVEN_VOICE_{gender.upper()}", ELEVEN_DEFAULT[gender])
         _eleven_tts(text, vid, out)
         return
     voice = pick_edge_voice(lang, gender)
@@ -190,3 +212,47 @@ def speak(text: str, lang: str, out: Path, gender: str = "Female", engine: str =
 
 def clean(text: str) -> str:
     return re.sub(r"\s+", " ", text).strip()
+
+
+# ------------------------------------------------------------------ lip-sync (GPU, pay-per-video)
+def lipsync_available() -> bool:
+    return bool(os.environ.get("REPLICATE_API_TOKEN"))
+
+
+def lipsync(video: Path, audio: Path, out: Path, progress=lambda m: None) -> None:
+    """Re-render mouths to match the new voice with ByteDance LatentSync on Replicate (~$0.10/run).
+    Model can be swapped with OXT_LIPSYNC_MODEL (owner/name)."""
+    import time as _t
+    import requests
+    tok = os.environ["REPLICATE_API_TOKEN"]
+    h = {"Authorization": f"Bearer {tok}"}
+    api = "https://api.replicate.com/v1"
+    model = os.environ.get("OXT_LIPSYNC_MODEL", "bytedance/latentsync")
+
+    def upload(p: Path) -> str:
+        with open(p, "rb") as f:
+            r = requests.post(f"{api}/files", headers=h, files={"content": (p.name, f)}, timeout=600)
+        r.raise_for_status()
+        return r.json()["urls"]["get"]
+
+    version = requests.get(f"{api}/models/{model}", headers=h, timeout=60).json()["latest_version"]["id"]
+    vurl, aurl = upload(video), upload(audio)
+    r = requests.post(f"{api}/predictions", headers={**h, "Content-Type": "application/json"},
+                      json={"version": version, "input": {"video": vurl, "audio": aurl}}, timeout=60)
+    r.raise_for_status()
+    pred = r.json()
+    t0 = _t.time()
+    while pred["status"] not in ("succeeded", "failed", "canceled"):
+        if _t.time() - t0 > 1800:
+            raise RuntimeError("Lip-sync took too long.")
+        progress(f"Lip-syncing on GPU ({int(_t.time() - t0)}s)")
+        _t.sleep(5)
+        pred = requests.get(pred["urls"]["get"], headers=h, timeout=60).json()
+    if pred["status"] != "succeeded":
+        raise RuntimeError(f"Lip-sync failed: {pred.get('error')}")
+    url = pred["output"] if isinstance(pred["output"], str) else pred["output"][0]
+    with requests.get(url, stream=True, timeout=600) as d:
+        d.raise_for_status()
+        with open(out, "wb") as f:
+            for chunk in d.iter_content(1 << 20):
+                f.write(chunk)

app/feedback.py

@@ -55,6 +55,18 @@ def voter_id(ip: str) -> str:
 
 
 def _learn(rec: dict) -> None:
+    if rec.get("kind") == "pair":  # contributed sentence pair counts as a vote both ways
+        if rec.get("dialect"):  # dialects are kept for training only, not mixed into the parent language
+            return
+        a, b = rec["en"].strip(), rec["text"].strip()
+        lang = rec["lang"]
+        if a and b:
+            _votes[_key("en", lang, a)][b].add(rec.get("voter", "?"))
+            _votes[_key(lang, "en", b)][a].add(rec.get("voter", "?"))
+            if rec.get("trusted"):
+                _trusted[_key("en", lang, a)] = b
+                _trusted[_key(lang, "en", b)] = a
+        return
     fix = (rec.get("correction") or "").strip()
     if not fix:
         return
@@ -91,7 +103,7 @@ def start() -> None:
             from huggingface_hub import CommitScheduler
             _scheduler = CommitScheduler(repo_id=REPO, repo_type="dataset", folder_path=FB_DIR,
                                          path_in_repo=".", every=5, private=True,
-                                         allow_patterns=["*.jsonl"])
+                                         allow_patterns=["*.jsonl", "audio/*"])
         except Exception as e:  # noqa: BLE001
             _sync_error = f"{type(e).__name__}: {str(e)[:200]}"
             print("feedback sync disabled:", e)
@@ -126,19 +138,41 @@ def examples(src: str, tgt: str, n: int = 15) -> list[dict]:
     return out[-n:]
 
 
+AUDIO_DIR = FB_DIR / "audio"
+AUDIO_DIR.mkdir(exist_ok=True)
+
+
+def add_pair(lang: str, en: str, text: str, voter: str, trusted: bool, dialect: str = "",
+             audio: bytes | None = None, audio_ext: str = "webm") -> None:
+    rec = {"kind": "pair", "lang": lang, "dialect": dialect, "en": en, "text": text,
+           "voter": voter, "trusted": trusted}
+    if audio:
+        name = f"{uuid.uuid4().hex[:12]}.{audio_ext}"
+        lock = _scheduler.lock if _scheduler else _lock
+        with lock:
+            (AUDIO_DIR / name).write_bytes(audio)
+        rec["audio"] = f"audio/{name}"
+    add(rec)
+
+
 def stats() -> dict:
-    n = good = bad = fixes = 0
+    n = good = bad = fixes = pairs = 0
+    by_lang: dict[str, int] = defaultdict(int)
     for p in FB_DIR.glob("*.jsonl"):
         for line in p.read_text(encoding="utf-8").splitlines():
             try:
                 r = json.loads(line)
             except ValueError:
                 continue
+            if r.get("kind") == "pair":
+                pairs += 1
+                by_lang[r.get("dialect") or r.get("lang")] += 1
+                continue
             n += 1
             good += r.get("rating") == 1
             bad += r.get("rating") == -1
             fixes += bool(r.get("correction"))
-    return {"ratings": n, "thumbs_up": good, "thumbs_down": bad, "corrections": fixes,
+    return {"ratings": n, "sentence_pairs": pairs, "pairs_by_language": dict(by_lang), "thumbs_up": good, "thumbs_down": bad, "corrections": fixes,
             "approved_fixes": len(_trusted) + sum(1 for k in _votes if lookup(*k)),
             "synced_to": REPO or None, "sync_active": _scheduler is not None,
             "sync_error": _sync_error or None}

app/pipeline.py

@@ -39,7 +39,10 @@ def pcm(path: Path, rate: int, tempo: float = 1.0) -> np.ndarray:
     return np.frombuffer(raw, np.int16).astype(np.float32)
 
 
-def fetch_url(url: str, folder: Path) -> Path:
+MEDIA_EXT = r"(?:mp4|m4v|mov|webm|mkv|m3u8|mp3|m4a|aac|wav|ogg|oga|opus|flac|amr|3gp)"
+
+
+def _ydl(url: str, folder: Path) -> Path:
     import yt_dlp
     opts = {"outtmpl": str(folder / "source.%(ext)s"), "quiet": True, "no_warnings": True,
             "noprogress": True, "noplaylist": True,
@@ -48,9 +51,42 @@ def fetch_url(url: str, folder: Path) -> Path:
     with yt_dlp.YoutubeDL(opts) as ydl:
         info = ydl.extract_info(url, download=True)
         p = Path(ydl.prepare_filename(info))
-    if not p.exists():
-        p = next(folder.glob("source.*"))
-    return p
+    return p if p.exists() else next(folder.glob("source.*"))
+
+
+def _find_media_in_page(url: str) -> list[str]:
+    """For ordinary websites: look for <video>/<audio> sources, og:video tags and media file links."""
+    import html
+    import re as _re
+    from urllib.parse import urljoin
+    import requests
+    page = requests.get(url, timeout=30, headers={"User-Agent": "Mozilla/5.0 SwitchLingo"}).text
+    found = []
+    pats = [r'<(?:video|audio|source)[^>]+src=["\']([^"\']+)',
+            r'<meta[^>]+(?:property|name)=["\'](?:og:video(?::secure_url|:url)?|og:audio|twitter:player:stream)["\'][^>]+content=["\']([^"\']+)',
+            r'["\'(]((?:https?:)?//[^"\'()\s]+?\.' + MEDIA_EXT + r'(?:\?[^"\'()\s]*)?)["\')]',
+            r'(?:youtube\.com/embed/|player\.vimeo\.com/video/)[\w-]+']
+    for pat in pats:
+        for m in _re.finditer(pat, page, _re.I):
+            u = html.unescape(m.group(1) if m.groups() else "https://" + m.group(0))
+            u = urljoin(url, u)
+            if u not in found:
+                found.append(u)
+    return found
+
+
+def fetch_url(url: str, folder: Path) -> Path:
+    """Social media and 1000+ video sites via yt-dlp; direct file links; any web page with a video."""
+    try:
+        return _ydl(url, folder)
+    except Exception as first:  # noqa: BLE001
+        for cand in _find_media_in_page(url)[:5]:
+            try:
+                return _ydl(cand, folder)
+            except Exception:  # noqa: BLE001
+                continue
+        raise RuntimeError("Couldn't find a video or audio on that link. If it's private or behind a "
+                           "login, download it and upload the file instead.") from first
 
 
 # ------------------------------------------------------------------ transcription
@@ -135,22 +171,78 @@ def burn_subs(video: Path, srt: Path, out: Path) -> None:
                     "-c:a", "aac", "-b:a", "160k", "-movflags", "+faststart", str(out)], check=True)
 
 
+def _speech_sample(media: Path, segs: list[dict], out: Path, max_sec: float = 75) -> Path:
+    """The speaker's own voice only (no silences), for voice cloning."""
+    keep, total = [], 0.0
+    for s in segs:
+        if total >= max_sec:
+            break
+        keep.append(f"between(t,{s['s']:.2f},{s['e']:.2f})")
+        total += s["e"] - s["s"]
+    subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", str(media), "-vn",
+                    "-af", f"aselect='{'+'.join(keep)}',asetpts=N/SR/TB", "-ac", "1",
+                    "-b:a", "128k", str(out)], check=True)
+    return out
+
+
 def dub(media: Path, segs: list[dict], lang: str, out: Path, has_video: bool, duration: float,
-        gender: str, voice_engine: str, work: Path, progress: Progress) -> None:
+        gender: str, voice_engine: str, work: Path, progress: Progress, clone: bool = False) -> None:
     parts = work / "voice"
     parts.mkdir(exist_ok=True)
+    n = max(len(segs), 1)
+
+    # 1) speak every line (ElevenLabs can clone the speaker directly)
+    voice_id = None
+    if clone and voice_engine == "elevenlabs":
+        progress("Cloning the speaker's voice", 0.70)
+        voice_id = engines.eleven_clone(_speech_sample(media, segs, work / "sample.mp3"))
+    try:
+        files = []
+        for k, s in enumerate(segs):
+            progress(f"Voicing line {k + 1} of {len(segs)}", 0.70 + 0.15 * k / n)
+            f = parts / f"{k:04}.mp3"
+            engines.speak(s["tr"], lang, f, gender=gender, engine=voice_engine, voice_id=voice_id)
+            files.append(f)
+    finally:
+        if voice_id:
+            engines.eleven_delete(voice_id)
+    clips = [pcm(f, SR) for f in files]
+
+    # 2) free cloning for every language: re-color each line to sound like the original speaker
+    if clone and voice_engine != "elevenlabs":
+        from . import voiceclone as vc
+        progress("Matching the speaker's voice", 0.86)
+        orig = pcm(media, vc.SR) / 32768.0
+        tgts = vc.speaker_embeddings(orig, [(s["s"], s["e"]) for s in segs])
+        tts22 = [pcm(f, vc.SR) / 32768.0 for f in files]
+        src_se = vc.embed(np.concatenate(tts22[:12]))
+        for k, (y, t) in enumerate(zip(tts22, tgts)):
+            progress(f"Matching voice, line {k + 1} of {len(segs)}", 0.86 + 0.09 * k / n)
+            if t is None:
+                continue
+            conv = vc.convert(y.astype(np.float32), src_se, t)
+            raw = work / "vc.raw"
+            (np.clip(conv, -1, 1) * 32767).astype(np.int16).tofile(raw)
+            clips[k] = np.frombuffer(subprocess.run(
+                ["ffmpeg", "-loglevel", "error", "-f", "s16le", "-ar", str(vc.SR), "-ac", "1", "-i", str(raw),
+                 "-f", "s16le", "-ar", str(SR), "-ac", "1", "-"], capture_output=True, check=True).stdout,
+                np.int16).astype(np.float32)
+
+    # 3) place lines at the speaker's timing (speed up max 1.2x, never overlap)
     track = np.zeros(int((duration + 15) * SR), np.float32)
     prev_end = 0
-    for k, s in enumerate(segs):
-        progress(f"Voicing line {k + 1} of {len(segs)}", 0.70 + 0.25 * k / max(len(segs), 1))
-        f = parts / f"{k:04}.mp3"
-        engines.speak(s["tr"], lang, f, gender=gender, engine=voice_engine)
+    for k, (s, a) in enumerate(zip(segs, clips)):
         nxt = segs[k + 1]["s"] if k + 1 < len(segs) else duration
         window = max(nxt - s["s"] - 0.05, 0.4)
-        a = pcm(f, SR)
-        if len(a) / SR > window:  # speed up a little to stay in sync (max 1.2x)
-            a = pcm(f, SR, min(len(a) / SR / window, 1.2))
-        st = max(int(s["s"] * SR), prev_end + int(0.08 * SR))  # never overlap
+        if len(a) / SR > window:
+            tempo = min(len(a) / SR / window, 1.2)
+            raw = work / "line.raw"
+            np.clip(a, -32768, 32767).astype(np.int16).tofile(raw)
+            a = np.frombuffer(subprocess.run(
+                ["ffmpeg", "-loglevel", "error", "-f", "s16le", "-ar", str(SR), "-ac", "1", "-i", str(raw),
+                 "-af", f"atempo={tempo:.3f}", "-f", "s16le", "-"], capture_output=True, check=True).stdout,
+                np.int16).astype(np.float32)
+        st = max(int(s["s"] * SR), prev_end + int(0.08 * SR))
         prev_end = st + len(a)
         track[st:st + len(a)] += a[:len(track) - st]
     raw = work / "dub.raw"
@@ -171,7 +263,8 @@ def dub(media: Path, segs: list[dict], lang: str, out: Path, has_video: bool, du
 def run(job_dir: Path, source_path: Path | None, url: str | None, target: str,
         source_lang: str | None = None, engine: str = "nllb", outputs: tuple = ("text", "subs", "dub"),
         quality: str = "fast", gender: str = "Female", voice_engine: str = "edge",
-        max_minutes: float = 0, progress: Progress = lambda m, p: None) -> dict:
+        max_minutes: float = 0, progress: Progress = lambda m, p: None,
+        clone: bool = True, lipsync: bool = False) -> dict:
     languages.get(target)
     progress("Getting the media", 0.02)
     media = fetch_url(url, job_dir) if url else source_path
@@ -220,12 +313,30 @@ def run(job_dir: Path, source_path: Path | None, url: str | None, target: str,
         progress("Creating the voice-over", 0.70)
         name = "voiceover.mp4" if meta["has_video"] else "voiceover.m4a"
         dub(base, segs, target, job_dir / name, meta["has_video"], meta["duration"],
-            gender, voice_engine, job_dir, progress)
+            gender, voice_engine, job_dir, progress, clone=clone)
         results["files"]["voiceover"] = name
+        results["voice_cloned"] = clone
+        if lipsync and meta["has_video"] and engines.lipsync_available():
+            progress("Lip-syncing mouths to the new voice (GPU)", 0.96)
+            vo = job_dir / name
+            aud = job_dir / "dub_audio.m4a"
+            subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", str(vo), "-vn", "-c:a", "copy", str(aud)],
+                           check=True)
+            synced = job_dir / "lipsync_raw.mp4"
+            engines.lipsync(vo, aud, synced, progress=lambda m: progress(m, 0.97))
+            subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", str(synced), "-i", str(aud),
+                            "-map", "0:v:0", "-map", "1:a:0", "-c:v", "libx264", "-preset", "veryfast",
+                            "-crf", "21", "-c:a", "copy", "-shortest", "-movflags", "+faststart",
+                            str(job_dir / "lipsync.mp4")], check=True)
+            results["files"]["lipsync"] = "lipsync.mp4"
+            for junk in (aud, synced):
+                junk.unlink(missing_ok=True)
 
     if meta["has_video"] or url:
         results["files"]["original"] = media.name
     shutil.rmtree(job_dir / "voice", ignore_errors=True)
+    for junk in ("sample.mp3", "vc.raw", "line.raw"):
+        (job_dir / junk).unlink(missing_ok=True)
     (job_dir / "dub.raw").unlink(missing_ok=True)
     progress("Done", 1.0)
     return results

app/server.py

@@ -83,7 +83,7 @@ def get_languages():
     return {"languages": languages.public_list(), "popular": languages.POPULAR_US,
             "limits": {"max_minutes": MAX_MINUTES, "max_mb": MAX_MB},
             "translators": engines.available_translators(),
-            "voices": engines.available_voices()}
+            "voices": engines.available_voices(), "lipsync": engines.lipsync_available()}
 
 
 @app.post("/api/jobs")
@@ -97,6 +97,8 @@ async def create_job(
     quality: str = Form("fast"),
     gender: str = Form("Female"),
     voice_engine: str = Form("edge"),
+    clone: str = Form("1"),
+    lipsync: str = Form("0"),
     file: UploadFile | None = File(None),
 ):
     if target not in languages.LANGS:
@@ -137,7 +139,8 @@ async def create_job(
     kwargs = dict(source_path=src_path, url=url.strip() or None, target=target,
                   source_lang=source or None, engine=engine,
                   outputs=tuple(o for o in outputs.split(",") if o), quality=quality,
-                  gender=gender, voice_engine=voice_engine, max_minutes=MAX_MINUTES)
+                  gender=gender, voice_engine=voice_engine, max_minutes=MAX_MINUTES,
+                  clone=clone == "1", lipsync=lipsync == "1" and engines.lipsync_available())
     pool.submit(_work, job_id, kwargs)
     return {"id": job_id}
 
@@ -206,6 +209,45 @@ def post_feedback(request: Request, job_id: str = Form(...), seg: int = Form(...
     return {"ok": True}
 
 
+TEACH_PROMPTS = [
+    "Where is the nearest hospital?", "How much does this cost?", "Thank you for helping my family.",
+    "What time does the store close?", "I don't understand, please speak slowly.",
+    "My child is sick and needs a doctor.", "Can you help me fill out this form?",
+    "The food is delicious.", "I am looking for work.", "Please call me back tomorrow.",
+    "Grandma is coming to visit next week.", "Where do I pay my bill?",
+    "I need a translator for my appointment.", "Happy New Year to everyone!",
+    "Be careful, the road is slippery.", "How do you say this in your language?",
+]
+TEACH_LANGS = {"lo": "Lao", "lo-isan": "Isan (Northeast Thai)", "hmn": "Hmong (White)",
+               "hmn-green": "Hmong (Green)", "km": "Khmer", "so": "Somali", "th": "Thai", "vi": "Vietnamese"}
+
+
+@app.get("/api/teach/languages")
+def teach_languages():
+    return {"languages": TEACH_LANGS, "prompts": TEACH_PROMPTS}
+
+
+@app.post("/api/teach")
+async def teach(request: Request, lang: str = Form(...), en: str = Form(...), text: str = Form(...),
+                token: str = Form(""), audio: UploadFile | None = File(None)):
+    if lang not in TEACH_LANGS:
+        raise HTTPException(400, "Unknown language")
+    en, text = en.strip()[:500], text.strip()[:500]
+    if not en or not text:
+        raise HTTPException(400, "Both sentences are needed")
+    data, ext = None, "webm"
+    if audio and audio.filename:
+        data = await audio.read(3 * 1024 * 1024 + 1)
+        if len(data) > 3 * 1024 * 1024:
+            raise HTTPException(413, "Recording is too long (max ~1 minute)")
+        ext = (Path(audio.filename).suffix.lstrip(".") or "webm")[:5]
+    base = lang.split("-")[0]  # dialects feed the parent language's memory
+    feedback.add_pair(base, en, text, feedback.voter_id(_client_ip(request)),
+                      bool(feedback.ADMIN_TOKEN) and token == feedback.ADMIN_TOKEN,
+                      dialect=TEACH_LANGS[lang] if "-" in lang else "", audio=data, audio_ext=ext)
+    return {"ok": True}
+
+
 @app.get("/api/feedback/stats")
 def feedback_stats():
     return feedback.stats()

app/static/index.html

@@ -5,6 +5,10 @@
 <meta name="viewport" content="width=device-width, initial-scale=1">
 <title>SwitchLingo — Any video, your language</title>
 <meta name="description" content="SwitchLingo: AI that translates any video or link into your language with subtitles and voice-over in 100+ languages, and gets better from your corrections.">
+<link rel="manifest" href="manifest.webmanifest">
+<meta name="theme-color" content="#1d1628">
+<link rel="apple-touch-icon" href="icon-192.png">
+<link rel="icon" href="icon-192.png">
 <link rel="preconnect" href="https://fonts.googleapis.com">
 <link href="https://fonts.googleapis.com/css2?family=Fraunces:opsz,wght@9..144,600&family=Inter:wght@400;500;600&display=swap" rel="stylesheet">
 <style>
@@ -55,11 +59,30 @@ td.t{white-space:nowrap;color:var(--muted);font-variant-numeric:tabular-nums;wid
 .fix{display:flex;gap:6px;margin-top:6px}.fix textarea{min-height:44px;font-size:14px}.fix button{border:0;background:var(--accent);color:var(--accent-ink);border-radius:8px;padding:0 12px;font-weight:600;cursor:pointer}
 .human{display:inline-block;font-size:11px;background:var(--soft);color:var(--ok);border-radius:99px;padding:1px 8px;margin-left:6px}
 .ok{color:var(--ok);font-size:12px}
+.sw{display:flex;align-items:center;justify-content:space-between;gap:12px;background:var(--soft);padding:12px 14px;border-radius:12px;margin-top:10px;cursor:pointer}
+.sw b{display:block;font-weight:600}.sw small{color:var(--muted)}
+.sw input{appearance:none;width:46px;height:28px;border-radius:99px;background:var(--line);position:relative;flex:none;cursor:pointer;transition:.2s}
+.sw input:checked{background:var(--accent)}
+.sw input::after{content:"";position:absolute;top:3px;left:3px;width:22px;height:22px;border-radius:50%;background:#fff;transition:.2s}
+.sw input:checked::after{left:21px}
+.badge{display:inline-block;font-size:11px;font-weight:600;background:var(--accent);color:var(--accent-ink);border-radius:99px;padding:1px 8px;margin-left:6px;vertical-align:middle}
+.rec{background:var(--soft);color:var(--ink)}.rec.on{background:var(--bad);color:#fff}
+.prompt{font:600 18px Fraunces,Georgia,serif;margin:6px 0 10px}
+#install{display:none}
+.sw.hidden{display:none}
+@media (max-width:640px){
+  .wrap{padding:18px 12px 90px}h1{font-size:30px}.card{padding:16px;border-radius:14px}
+  .btn{width:100%;justify-content:center;padding:15px}.btn.ghost{width:auto}
+  input[type=text],input[type=url],select,textarea{font-size:16px;padding:13px}
+  table thead{display:none}table,tbody,tr,td{display:block;width:100%}
+  tr{border-bottom:1px solid var(--line);padding:10px 0}td{border:0;padding:3px 0}
+  td.t{width:auto}.fb{width:auto}.fb button{padding:8px 14px;font-size:17px;margin:6px 6px 0 0}
+}
 </style>
 </head>
 <body>
 <div class="wrap">
-<header><h1>Switch<span>Lingo</span></h1><div class="tag">AI video translation. Subtitles and voice-over in 100+ languages, and it learns from your fixes.</div></header>
+<header><h1>Switch<span>Lingo</span></h1><button id="install" class="btn ghost" style="margin:0 0 0 auto">📲 Add to phone</button><div class="tag">AI video translation. Subtitles and voice-over in 100+ languages, and it learns from your fixes.</div></header>
 
 <section class="card demo">
   <h2>See it switch</h2>
@@ -69,11 +92,15 @@ td.t{white-space:nowrap;color:var(--muted);font-variant-numeric:tabular-nums;wid
 
 <section class="card">
   <h2>1. Add your video or audio</h2>
-  <div class="tabs"><button id="tLink" class="on">Paste a link</button><button id="tFile">Upload a file</button></div>
-  <div id="pLink"><input id="url" type="url" placeholder="Facebook, Instagram, YouTube, TikTok… link"></div>
+  <div class="tabs"><button id="tLink" class="on">🔗 Link</button><button id="tFile">📁 Upload</button><button id="tRec">🎤 Record</button></div>
+  <div id="pLink"><input id="url" type="url" inputmode="url" placeholder="Facebook, Instagram, TikTok, YouTube, X, or any website with a video"></div>
   <div id="pFile" class="hidden">
-    <div id="drop" class="drop">Drop a video/audio file here or <u>browse</u><br><small id="fname"></small></div>
-    <input id="file" type="file" accept="video/*,audio/*" hidden>
+    <div id="drop" class="drop">Tap to choose a video or voice recording<br><small class="muted">mp4, mov, m4a, mp3, wav, amr, 3gp, ogg, opus…</small><br><small id="fname"></small></div>
+    <input id="file" type="file" accept="video/*,audio/*,.m4a,.mp3,.wav,.amr,.3gp,.3ga,.aac,.ogg,.oga,.opus,.flac,.wma,.webm,.mkv,.mov,.mp4" hidden>
+  </div>
+  <div id="pRec" class="hidden">
+    <div class="drop" id="recMain">🎤 Tap to start recording<br><small class="muted">Speak or hold your phone near the speaker. Up to 10 minutes.</small></div>
+    <audio id="recMainPlay" controls class="hidden"></audio>
   </div>
 
   <div class="row">
@@ -86,6 +113,8 @@ td.t{white-space:nowrap;color:var(--muted);font-variant-numeric:tabular-nums;wid
     <label class="chk"><input type="checkbox" class="out" value="subs" checked> Subtitles on video</label>
     <label class="chk"><input type="checkbox" class="out" value="dub" checked> Voice-over</label>
   </div>
+  <label class="sw"><span><b>🎙️ Keep the speaker's voice</b><small>Voice-over sounds like the person in the video</small></span><input type="checkbox" id="clone" checked></label>
+  <label class="sw hidden" id="lipWrap"><span><b>👄 Lip-sync <span class="badge">PRO</span></b><small>Mouths move with the new language (takes ~2 min longer)</small></span><input type="checkbox" id="lipsync"></label>
 
   <details>
     <summary>Options</summary>
@@ -98,7 +127,7 @@ td.t{white-space:nowrap;color:var(--muted);font-variant-numeric:tabular-nums;wid
     </div>
   </details>
   <button id="go" class="btn">Translate</button>
-  <p class="muted" style="font-size:13px;margin:10px 0 0"><span id="limits"></span> Only use videos you have the right to share. Voice-overs are AI voices, not the original speaker.</p>
+  <p class="muted" style="font-size:13px;margin:10px 0 0"><span id="limits"></span> Only use videos you have the right to share. AI voices, even cloned ones, are never the real person speaking. Get permission before reposting.</p>
   <div id="formErr" class="err"></div>
 </section>
 
@@ -123,6 +152,22 @@ td.t{white-space:nowrap;color:var(--muted);font-variant-numeric:tabular-nums;wid
   <div id="jobs" class="jobs muted">Nothing yet.</div>
 </section>
 
+<section class="card" id="teach">
+  <h2>🇱🇦 Teach SwitchLingo your language</h2>
+  <p class="muted" style="margin-top:-6px">Big apps barely know Lao, Isan or Hmong. Every sentence you add makes SwitchLingo better for your family and community. Takes 20 seconds.</p>
+  <form id="teachForm">
+    <label class="lb" for="tlang">Language</label><select id="tlang" name="lang"></select>
+    <div style="margin-top:12px"><label class="lb">Say this in your language <a href="#" id="nextP" style="float:right">another one ↻</a></label>
+      <div class="prompt" id="tprompt"></div><input type="hidden" name="en" id="ten"></div>
+    <label class="lb" for="ttext">Type it (your own spelling is fine)</label><textarea id="ttext" name="text" required></textarea>
+    <div style="display:flex;gap:8px;align-items:center;flex-wrap:wrap;margin-top:10px">
+      <button type="button" class="btn ghost rec" id="recBtn">🎤 Record yourself saying it</button><audio id="recPlay" controls class="hidden" style="max-width:260px;margin:0"></audio>
+    </div>
+    <button class="btn" type="submit">Add my sentence</button> <span id="teachMsg" class="muted"></span>
+  </form>
+  <p class="muted" style="font-size:13px" id="teachStats"></p>
+</section>
+
 <section class="card">
   <h2>Don't see your language?</h2>
   <p class="muted" style="margin-top:-6px">Tell us what to add. Requests help decide which languages and dialects get trained next.</p>
@@ -143,9 +188,22 @@ let mode = "link", file = null, LANG = {};
 const esc = s => (s ?? "").replace(/[&<>"]/g, c => ({"&":"&amp;","<":"&lt;",">":"&gt;",'"':"&quot;"}[c]));
 const fmt = s => { s = Math.floor(s); return `${Math.floor(s/60)}:${String(s%60).padStart(2,"0")}`; };
 
-function setMode(m){ mode=m; $("#tLink").classList.toggle("on",m==="link"); $("#tFile").classList.toggle("on",m==="file");
-  $("#pLink").classList.toggle("hidden",m!=="link"); $("#pFile").classList.toggle("hidden",m!=="file"); }
-$("#tLink").onclick=()=>setMode("link"); $("#tFile").onclick=()=>setMode("file");
+function setMode(m){ mode=m; [["link","Link"],["file","File"],["rec","Rec"]].forEach(([k,n])=>{
+  $("#t"+n).classList.toggle("on",m===k); $("#p"+n).classList.toggle("hidden",m!==k); }); }
+$("#tLink").onclick=()=>setMode("link"); $("#tFile").onclick=()=>setMode("file"); $("#tRec").onclick=()=>setMode("rec");
+let mainRec=null, mainBlob=null, recTimer=null;
+$("#recMain").onclick=async()=>{
+  if(mainRec&&mainRec.state==="recording"){ mainRec.stop(); return; }
+  try{
+    const st=await navigator.mediaDevices.getUserMedia({audio:true}); const ch=[]; const t0=Date.now();
+    mainRec=new MediaRecorder(st); mainRec.ondataavailable=e=>ch.push(e.data);
+    mainRec.onstop=()=>{ clearInterval(recTimer); st.getTracks().forEach(t=>t.stop()); mainBlob=new Blob(ch,{type:mainRec.mimeType});
+      $("#recMainPlay").src=URL.createObjectURL(mainBlob); $("#recMainPlay").classList.remove("hidden");
+      $("#recMain").innerHTML="✅ Recorded. Tap to record again"; };
+    mainRec.start(); recTimer=setInterval(()=>{ const s=Math.floor((Date.now()-t0)/1000);
+      $("#recMain").innerHTML=`⏹ Recording ${fmt(s)}… tap to stop`; if(s>=600) mainRec.stop(); },500);
+  }catch(e){ $("#formErr").textContent="Microphone not allowed. Check your browser settings."; }
+};
 const drop=$("#drop");
 drop.onclick=()=>$("#file").click();
 $("#file").onchange=e=>{file=e.target.files[0]; $("#fname").textContent=file?file.name:"";};
@@ -169,6 +227,10 @@ async function init(){
   r.voices.forEach(v=>$("#voice").insertAdjacentHTML("beforeend",`<option value="${v}">${vn[v]||v}</option>`));
   if(r.voices.includes("elevenlabs")) $("#voice").value="elevenlabs";
   if(r.translators.includes("claude")) $("#engine").value="claude";
+  if(r.lipsync) $("#lipWrap").classList.remove("hidden");
+  const qs=new URLSearchParams(location.search), shared=qs.get("url")||qs.get("text")||"";
+  const m=shared.match(/https?:\/\/\S+/); if(m){ setMode("link"); $("#url").value=m[0]; $("#go").scrollIntoView({behavior:"smooth",block:"center"}); }
+  initTeach();
   loadJobs();
   const id=new URLSearchParams(location.search).get("job"); if(id) watch(id);
 }
@@ -177,10 +239,12 @@ $("#go").onclick=async()=>{
   $("#formErr").textContent="";
   const fd=new FormData();
   if(mode==="link"){ if(!$("#url").value.trim()) return $("#formErr").textContent="Paste a link first."; fd.append("url",$("#url").value.trim()); }
-  else { if(!file) return $("#formErr").textContent="Choose a file first."; fd.append("file",file); }
+  else if(mode==="file"){ if(!file) return $("#formErr").textContent="Choose a file first."; fd.append("file",file); }
+  else { if(!mainBlob) return $("#formErr").textContent="Record something first."; fd.append("file",mainBlob,"recording."+((mainBlob.type.split("/")[1]||"webm").split(";")[0])); }
   const outs=["text",...[...document.querySelectorAll(".out:checked")].map(c=>c.value)];
   Object.entries({target:$("#tgt").value,source:$("#src").value,engine:$("#engine").value,outputs:outs.join(","),
-    quality:$("#quality").value,gender:$("#gender").value,voice_engine:$("#voice").value}).forEach(([k,v])=>fd.append(k,v));
+    quality:$("#quality").value,gender:$("#gender").value,voice_engine:$("#voice").value,
+    clone:$("#clone").checked?"1":"0",lipsync:$("#lipsync").checked?"1":"0"}).forEach(([k,v])=>fd.append(k,v));
   $("#go").disabled=true; $("#go").textContent="Uploading…";
   try{
     const r=await fetch("/api/jobs",{method:"POST",body:fd}); const j=await r.json();
@@ -205,10 +269,11 @@ async function watch(id){
 function show(j){
   const r=j.result, f=r.files, url=n=>`/api/files/${j.id}/${encodeURIComponent(n)}`;
   $("#langs").innerHTML=`<span class="pill">${esc(r.source_name)} → ${esc(r.target_name)}</span><span class="pill">${fmt(r.duration)}</span>`
+    + (r.voice_cloned?'<span class="pill">🎙️ speaker\'s voice</span>':'') + (f.lipsync?'<span class="pill">👄 lip-synced</span>':'')
     + (r.auto_detected && r.detect_confidence < 0.85 ? `<p class="err">Heard ${esc(r.source_name)}, but only ${Math.round(r.detect_confidence*100)}% sure. If that's wrong, set "Spoken language" above and translate again. Mixed-language speech (e.g. Somali with English words) often fools auto-detect.</p>` : "");
-  const main=f.voiceover||f.subtitled;
+  const main=f.lipsync||f.voiceover||f.subtitled;
   $("#player").innerHTML = main ? (main.endsWith(".m4a")?`<audio controls src="${url(main)}"></audio>`:`<video controls playsinline src="${url(main)}"></video>`) : "";
-  const labels={voiceover:"Voice-over",subtitled:"Subtitled video",translated_srt:"Translated subtitles (.srt)",original_srt:"Original subtitles (.srt)",transcript:"Side-by-side transcript (.txt)",original:"Original media"};
+  const labels={lipsync:"Lip-synced video",voiceover:"Voice-over",subtitled:"Subtitled video",translated_srt:"Translated subtitles (.srt)",original_srt:"Original subtitles (.srt)",transcript:"Side-by-side transcript (.txt)",original:"Original media"};
   $("#downloads").innerHTML=Object.entries(labels).filter(([k])=>f[k]).map(([k,l])=>`<a class="btn ghost" href="${url(f[k])}?download=true">${l}</a>`).join("");
   $("#thO").textContent="Before — "+r.source_name; $("#thT").textContent="After — "+r.target_name;
   CUR=j.id;
@@ -241,6 +306,41 @@ async function loadJobs(){
 $("#reqForm").onsubmit=async e=>{ e.preventDefault();
   const r=await fetch("/api/language-requests",{method:"POST",body:new FormData(e.target)});
   $("#reqMsg").textContent=r.ok?"Thanks, request saved.":"Couldn't save, try again."; if(r.ok) e.target.reset(); };
+let recBlob=null, recorder=null, PROMPTS=[];
+async function initTeach(){
+  const r=await (await fetch("/api/teach/languages")).json(); PROMPTS=r.prompts;
+  $("#tlang").innerHTML=Object.entries(r.languages).map(([k,v])=>`<option value="${k}">${esc(v)}</option>`).join("");
+  nextPrompt(); teachStats();
+}
+function nextPrompt(){ const p=PROMPTS[Math.floor(Math.random()*PROMPTS.length)]; $("#tprompt").textContent="“"+p+"”"; $("#ten").value=p; }
+$("#nextP").onclick=e=>{e.preventDefault(); nextPrompt();};
+async function teachStats(){ try{ const s=await (await fetch("/api/feedback/stats")).json();
+  const by=Object.entries(s.pairs_by_language||{}).map(([k,v])=>`${esc(k)} ${v}`).join(" · ");
+  $("#teachStats").textContent=`${s.sentence_pairs||0} sentences taught so far${by?" ("+by+")":""} · ${s.corrections||0} fixes from viewers.`; }catch(e){} }
+$("#recBtn").onclick=async()=>{
+  if(recorder&&recorder.state==="recording"){ recorder.stop(); return; }
+  try{
+    const st=await navigator.mediaDevices.getUserMedia({audio:true}); const ch=[];
+    recorder=new MediaRecorder(st); recorder.ondataavailable=e=>ch.push(e.data);
+    recorder.onstop=()=>{ st.getTracks().forEach(t=>t.stop()); recBlob=new Blob(ch,{type:recorder.mimeType});
+      $("#recPlay").src=URL.createObjectURL(recBlob); $("#recPlay").classList.remove("hidden");
+      $("#recBtn").classList.remove("on"); $("#recBtn").textContent="🎤 Record again"; };
+    recorder.start(); $("#recBtn").classList.add("on"); $("#recBtn").textContent="⏹ Stop recording";
+    setTimeout(()=>recorder&&recorder.state==="recording"&&recorder.stop(),30000);
+  }catch(e){ $("#teachMsg").textContent="Microphone not allowed. You can still type it."; }
+};
+$("#teachForm").onsubmit=async e=>{ e.preventDefault();
+  const fd=new FormData(e.target); fd.append("token",ADMIN);
+  if(recBlob) fd.append("audio",recBlob,"rec."+((recBlob.type.split("/")[1]||"webm").split(";")[0]));
+  const r=await fetch("/api/teach",{method:"POST",body:fd});
+  $("#teachMsg").textContent=r.ok?"Thank you! 🙏 Added.":"Couldn't save, try again.";
+  if(r.ok){ $("#ttext").value=""; recBlob=null; $("#recPlay").classList.add("hidden"); $("#recBtn").textContent="🎤 Record yourself saying it"; nextPrompt(); teachStats(); } };
+
+// phone: installable app + share-to-SwitchLingo
+if("serviceWorker" in navigator) navigator.serviceWorker.register("sw.js").catch(()=>{});
+let deferred=null;
+window.addEventListener("beforeinstallprompt",e=>{ e.preventDefault(); deferred=e; $("#install").style.display="inline-flex"; });
+$("#install").onclick=async()=>{ if(!deferred) return; deferred.prompt(); await deferred.userChoice; deferred=null; $("#install").style.display="none"; };
 init();
 </script>
 </body>

app/static/manifest.webmanifest

@@ -0,0 +1,19 @@
+{
+  "name": "SwitchLingo — any video, your language",
+  "short_name": "SwitchLingo",
+  "description": "Translate any video into your language with subtitles and a voice that sounds like the speaker.",
+  "start_url": "/",
+  "scope": "/",
+  "display": "standalone",
+  "background_color": "#1d1628",
+  "theme_color": "#1d1628",
+  "icons": [
+    {"src": "icon-192.png", "sizes": "192x192", "type": "image/png", "purpose": "any maskable"},
+    {"src": "icon-512.png", "sizes": "512x512", "type": "image/png", "purpose": "any maskable"}
+  ],
+  "share_target": {
+    "action": "/",
+    "method": "GET",
+    "params": {"title": "title", "text": "text", "url": "url"}
+  }
+}

app/static/sw.js

@@ -0,0 +1,4 @@
+// Minimal service worker: makes SwitchLingo installable; always uses the live site (no stale cache).
+self.addEventListener("install", () => self.skipWaiting());
+self.addEventListener("activate", e => e.waitUntil(self.clients.claim()));
+self.addEventListener("fetch", () => {});

app/voiceclone/__init__.py

@@ -0,0 +1,86 @@
+"""Free voice cloning for every language.
+
+The voice-over is first spoken by a normal neural voice in the target language, then OpenVoice V2's
+tone-color converter (MIT, runs on CPU) re-colors it so it sounds like the person in the original
+video. Because it converts audio instead of generating speech, it works for any language,
+including Lao, Somali and Khmer.
+
+Each line takes its voice from what the speaker sounded like at that moment, smoothed with nearby
+lines that sound like the same person, so videos with two or three speakers keep them apart.
+"""
+from __future__ import annotations
+
+import threading
+from functools import lru_cache
+
+import numpy as np
+
+SR = 22050
+_lock = threading.Lock()
+
+
+@lru_cache(maxsize=1)
+def _model():
+    import torch
+    from huggingface_hub import hf_hub_download
+    from .ov import utils
+    from .ov.models import SynthesizerTrn
+    cfg = hf_hub_download("myshell-ai/OpenVoiceV2", "converter/config.json")
+    ckpt = hf_hub_download("myshell-ai/OpenVoiceV2", "converter/checkpoint.pth")
+    hps = utils.get_hparams_from_file(cfg)
+    m = SynthesizerTrn(len(getattr(hps, "symbols", [])), hps.data.filter_length // 2 + 1,
+                       n_speakers=hps.data.n_speakers, **hps.model)
+    m.load_state_dict(torch.load(ckpt, map_location="cpu")["model"], strict=False)
+    m.eval()
+    torch.set_num_threads(max(1, min(8, torch.get_num_threads())))
+    return m, hps
+
+
+def _spec(y: np.ndarray):
+    import torch
+    from .ov.mel_processing import spectrogram_torch
+    _, hps = _model()
+    t = torch.FloatTensor(y).unsqueeze(0)
+    return spectrogram_torch(t, hps.data.filter_length, hps.data.sampling_rate, hps.data.hop_length,
+                             hps.data.win_length, center=False)
+
+
+def embed(y: np.ndarray):
+    """Speaker embedding of a mono float clip at 22.05 kHz."""
+    import torch
+    m, _ = _model()
+    with torch.no_grad():
+        return m.ref_enc(_spec(y).transpose(1, 2)).unsqueeze(-1)
+
+
+def convert(y: np.ndarray, src_se, tgt_se, tau: float = 0.3) -> np.ndarray:
+    import torch
+    m, _ = _model()
+    with _lock, torch.no_grad():
+        spec = _spec(y)
+        out = m.voice_conversion(spec, torch.LongTensor([spec.size(-1)]), sid_src=src_se,
+                                 sid_tgt=tgt_se, tau=tau)[0][0, 0]
+    return out.float().numpy()
+
+
+def speaker_embeddings(original: np.ndarray, spans: list[tuple[float, float]], min_len: float = 3.0):
+    """One target voice per line: that moment of the original, widened to >=3s, then averaged
+    with other lines that sound like the same speaker (cosine > 0.85)."""
+    import torch
+    raw = []
+    for s, e in spans:
+        if e - s < min_len:
+            pad = (min_len - (e - s)) / 2
+            s, e = max(0.0, s - pad), e + pad
+        clip = original[int(s * SR): int(e * SR)]
+        raw.append(embed(clip) if len(clip) > SR // 2 else None)
+    good = [r for r in raw if r is not None]
+    if not good:
+        return [None] * len(spans)
+    glob = torch.stack(good).mean(0)
+    raw = [r if r is not None else glob for r in raw]
+    flat = torch.stack([r.flatten() for r in raw])
+    flat = flat / flat.norm(dim=1, keepdim=True)
+    sim = flat @ flat.T
+    return [torch.stack([raw[j] for j in range(len(raw)) if sim[i, j] > 0.85]).mean(0)
+            for i in range(len(raw))]

app/voiceclone/ov/LICENSE

@@ -0,0 +1,7 @@
+Copyright 2024 MyShell.ai
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

app/voiceclone/ov/__init__.py

@@ -0,0 +1 @@
+"""Vendored from myshell-ai/OpenVoice (MIT). Only the tone-color converter parts."""

app/voiceclone/ov/attentions.py

@@ -0,0 +1,465 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from . import commons
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        x = x.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
+
+
+@torch.jit.script
+def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
+    n_channels_int = n_channels[0]
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :n_channels_int, :])
+    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+    acts = t_act * s_act
+    return acts
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size=1,
+        p_dropout=0.0,
+        window_size=4,
+        isflow=True,
+        **kwargs
+    ):
+        super().__init__()
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.window_size = window_size
+        # if isflow:
+        #  cond_layer = torch.nn.Conv1d(256, 2*hidden_channels*n_layers, 1)
+        #  self.cond_pre = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, 1)
+        #  self.cond_layer = weight_norm(cond_layer, name='weight')
+        #  self.gin_channels = 256
+        self.cond_layer_idx = self.n_layers
+        if "gin_channels" in kwargs:
+            self.gin_channels = kwargs["gin_channels"]
+            if self.gin_channels != 0:
+                self.spk_emb_linear = nn.Linear(self.gin_channels, self.hidden_channels)
+                # vits2 says 3rd block, so idx is 2 by default
+                self.cond_layer_idx = (
+                    kwargs["cond_layer_idx"] if "cond_layer_idx" in kwargs else 2
+                )
+                # logging.debug(self.gin_channels, self.cond_layer_idx)
+                assert (
+                    self.cond_layer_idx < self.n_layers
+                ), "cond_layer_idx should be less than n_layers"
+        self.drop = nn.Dropout(p_dropout)
+        self.attn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_2 = nn.ModuleList()
+
+        for i in range(self.n_layers):
+            self.attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels,
+                    hidden_channels,
+                    n_heads,
+                    p_dropout=p_dropout,
+                    window_size=window_size,
+                )
+            )
+            self.norm_layers_1.append(LayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FFN(
+                    hidden_channels,
+                    hidden_channels,
+                    filter_channels,
+                    kernel_size,
+                    p_dropout=p_dropout,
+                )
+            )
+            self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+    def forward(self, x, x_mask, g=None):
+        attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for i in range(self.n_layers):
+            if i == self.cond_layer_idx and g is not None:
+                g = self.spk_emb_linear(g.transpose(1, 2))
+                g = g.transpose(1, 2)
+                x = x + g
+                x = x * x_mask
+            y = self.attn_layers[i](x, x, attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_1[i](x + y)
+
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.drop(y)
+            x = self.norm_layers_2[i](x + y)
+        x = x * x_mask
+        return x
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size=1,
+        p_dropout=0.0,
+        proximal_bias=False,
+        proximal_init=True,
+        **kwargs
+    ):
+        super().__init__()
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
+
+        self.drop = nn.Dropout(p_dropout)
+        self.self_attn_layers = nn.ModuleList()
+        self.norm_layers_0 = nn.ModuleList()
+        self.encdec_attn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_2 = nn.ModuleList()
+        for i in range(self.n_layers):
+            self.self_attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels,
+                    hidden_channels,
+                    n_heads,
+                    p_dropout=p_dropout,
+                    proximal_bias=proximal_bias,
+                    proximal_init=proximal_init,
+                )
+            )
+            self.norm_layers_0.append(LayerNorm(hidden_channels))
+            self.encdec_attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout
+                )
+            )
+            self.norm_layers_1.append(LayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FFN(
+                    hidden_channels,
+                    hidden_channels,
+                    filter_channels,
+                    kernel_size,
+                    p_dropout=p_dropout,
+                    causal=True,
+                )
+            )
+            self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+    def forward(self, x, x_mask, h, h_mask):
+        """
+        x: decoder input
+        h: encoder output
+        """
+        self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(
+            device=x.device, dtype=x.dtype
+        )
+        encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for i in range(self.n_layers):
+            y = self.self_attn_layers[i](x, x, self_attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_0[i](x + y)
+
+            y = self.encdec_attn_layers[i](x, h, encdec_attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_1[i](x + y)
+
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.drop(y)
+            x = self.norm_layers_2[i](x + y)
+        x = x * x_mask
+        return x
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        channels,
+        out_channels,
+        n_heads,
+        p_dropout=0.0,
+        window_size=None,
+        heads_share=True,
+        block_length=None,
+        proximal_bias=False,
+        proximal_init=False,
+    ):
+        super().__init__()
+        assert channels % n_heads == 0
+
+        self.channels = channels
+        self.out_channels = out_channels
+        self.n_heads = n_heads
+        self.p_dropout = p_dropout
+        self.window_size = window_size
+        self.heads_share = heads_share
+        self.block_length = block_length
+        self.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
+        self.attn = None
+
+        self.k_channels = channels // n_heads
+        self.conv_q = nn.Conv1d(channels, channels, 1)
+        self.conv_k = nn.Conv1d(channels, channels, 1)
+        self.conv_v = nn.Conv1d(channels, channels, 1)
+        self.conv_o = nn.Conv1d(channels, out_channels, 1)
+        self.drop = nn.Dropout(p_dropout)
+
+        if window_size is not None:
+            n_heads_rel = 1 if heads_share else n_heads
+            rel_stddev = self.k_channels**-0.5
+            self.emb_rel_k = nn.Parameter(
+                torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
+                * rel_stddev
+            )
+            self.emb_rel_v = nn.Parameter(
+                torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
+                * rel_stddev
+            )
+
+        nn.init.xavier_uniform_(self.conv_q.weight)
+        nn.init.xavier_uniform_(self.conv_k.weight)
+        nn.init.xavier_uniform_(self.conv_v.weight)
+        if proximal_init:
+            with torch.no_grad():
+                self.conv_k.weight.copy_(self.conv_q.weight)
+                self.conv_k.bias.copy_(self.conv_q.bias)
+
+    def forward(self, x, c, attn_mask=None):
+        q = self.conv_q(x)
+        k = self.conv_k(c)
+        v = self.conv_v(c)
+
+        x, self.attn = self.attention(q, k, v, mask=attn_mask)
+
+        x = self.conv_o(x)
+        return x
+
+    def attention(self, query, key, value, mask=None):
+        # reshape [b, d, t] -> [b, n_h, t, d_k]
+        b, d, t_s, t_t = (*key.size(), query.size(2))
+        query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
+        key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+        value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+
+        scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1))
+        if self.window_size is not None:
+            assert (
+                t_s == t_t
+            ), "Relative attention is only available for self-attention."
+            key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
+            rel_logits = self._matmul_with_relative_keys(
+                query / math.sqrt(self.k_channels), key_relative_embeddings
+            )
+            scores_local = self._relative_position_to_absolute_position(rel_logits)
+            scores = scores + scores_local
+        if self.proximal_bias:
+            assert t_s == t_t, "Proximal bias is only available for self-attention."
+            scores = scores + self._attention_bias_proximal(t_s).to(
+                device=scores.device, dtype=scores.dtype
+            )
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, -1e4)
+            if self.block_length is not None:
+                assert (
+                    t_s == t_t
+                ), "Local attention is only available for self-attention."
+                block_mask = (
+                    torch.ones_like(scores)
+                    .triu(-self.block_length)
+                    .tril(self.block_length)
+                )
+                scores = scores.masked_fill(block_mask == 0, -1e4)
+        p_attn = F.softmax(scores, dim=-1)  # [b, n_h, t_t, t_s]
+        p_attn = self.drop(p_attn)
+        output = torch.matmul(p_attn, value)
+        if self.window_size is not None:
+            relative_weights = self._absolute_position_to_relative_position(p_attn)
+            value_relative_embeddings = self._get_relative_embeddings(
+                self.emb_rel_v, t_s
+            )
+            output = output + self._matmul_with_relative_values(
+                relative_weights, value_relative_embeddings
+            )
+        output = (
+            output.transpose(2, 3).contiguous().view(b, d, t_t)
+        )  # [b, n_h, t_t, d_k] -> [b, d, t_t]
+        return output, p_attn
+
+    def _matmul_with_relative_values(self, x, y):
+        """
+        x: [b, h, l, m]
+        y: [h or 1, m, d]
+        ret: [b, h, l, d]
+        """
+        ret = torch.matmul(x, y.unsqueeze(0))
+        return ret
+
+    def _matmul_with_relative_keys(self, x, y):
+        """
+        x: [b, h, l, d]
+        y: [h or 1, m, d]
+        ret: [b, h, l, m]
+        """
+        ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
+        return ret
+
+    def _get_relative_embeddings(self, relative_embeddings, length):
+        2 * self.window_size + 1
+        # Pad first before slice to avoid using cond ops.
+        pad_length = max(length - (self.window_size + 1), 0)
+        slice_start_position = max((self.window_size + 1) - length, 0)
+        slice_end_position = slice_start_position + 2 * length - 1
+        if pad_length > 0:
+            padded_relative_embeddings = F.pad(
+                relative_embeddings,
+                commons.convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]),
+            )
+        else:
+            padded_relative_embeddings = relative_embeddings
+        used_relative_embeddings = padded_relative_embeddings[
+            :, slice_start_position:slice_end_position
+        ]
+        return used_relative_embeddings
+
+    def _relative_position_to_absolute_position(self, x):
+        """
+        x: [b, h, l, 2*l-1]
+        ret: [b, h, l, l]
+        """
+        batch, heads, length, _ = x.size()
+        # Concat columns of pad to shift from relative to absolute indexing.
+        x = F.pad(x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, 1]]))
+
+        # Concat extra elements so to add up to shape (len+1, 2*len-1).
+        x_flat = x.view([batch, heads, length * 2 * length])
+        x_flat = F.pad(
+            x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [0, length - 1]])
+        )
+
+        # Reshape and slice out the padded elements.
+        x_final = x_flat.view([batch, heads, length + 1, 2 * length - 1])[
+            :, :, :length, length - 1 :
+        ]
+        return x_final
+
+    def _absolute_position_to_relative_position(self, x):
+        """
+        x: [b, h, l, l]
+        ret: [b, h, l, 2*l-1]
+        """
+        batch, heads, length, _ = x.size()
+        # pad along column
+        x = F.pad(
+            x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length - 1]])
+        )
+        x_flat = x.view([batch, heads, length**2 + length * (length - 1)])
+        # add 0's in the beginning that will skew the elements after reshape
+        x_flat = F.pad(x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [length, 0]]))
+        x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:]
+        return x_final
+
+    def _attention_bias_proximal(self, length):
+        """Bias for self-attention to encourage attention to close positions.
+        Args:
+          length: an integer scalar.
+        Returns:
+          a Tensor with shape [1, 1, length, length]
+        """
+        r = torch.arange(length, dtype=torch.float32)
+        diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
+        return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)
+
+
+class FFN(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        filter_channels,
+        kernel_size,
+        p_dropout=0.0,
+        activation=None,
+        causal=False,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.activation = activation
+        self.causal = causal
+
+        if causal:
+            self.padding = self._causal_padding
+        else:
+            self.padding = self._same_padding
+
+        self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size)
+        self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
+        self.drop = nn.Dropout(p_dropout)
+
+    def forward(self, x, x_mask):
+        x = self.conv_1(self.padding(x * x_mask))
+        if self.activation == "gelu":
+            x = x * torch.sigmoid(1.702 * x)
+        else:
+            x = torch.relu(x)
+        x = self.drop(x)
+        x = self.conv_2(self.padding(x * x_mask))
+        return x * x_mask
+
+    def _causal_padding(self, x):
+        if self.kernel_size == 1:
+            return x
+        pad_l = self.kernel_size - 1
+        pad_r = 0
+        padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, commons.convert_pad_shape(padding))
+        return x
+
+    def _same_padding(self, x):
+        if self.kernel_size == 1:
+            return x
+        pad_l = (self.kernel_size - 1) // 2
+        pad_r = self.kernel_size // 2
+        padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, commons.convert_pad_shape(padding))
+        return x

app/voiceclone/ov/commons.py

@@ -0,0 +1,160 @@
+import math
+import torch
+from torch.nn import functional as F
+
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+def convert_pad_shape(pad_shape):
+    layer = pad_shape[::-1]
+    pad_shape = [item for sublist in layer for item in sublist]
+    return pad_shape
+
+
+def intersperse(lst, item):
+    result = [item] * (len(lst) * 2 + 1)
+    result[1::2] = lst
+    return result
+
+
+def kl_divergence(m_p, logs_p, m_q, logs_q):
+    """KL(P||Q)"""
+    kl = (logs_q - logs_p) - 0.5
+    kl += (
+        0.5 * (torch.exp(2.0 * logs_p) + ((m_p - m_q) ** 2)) * torch.exp(-2.0 * logs_q)
+    )
+    return kl
+
+
+def rand_gumbel(shape):
+    """Sample from the Gumbel distribution, protect from overflows."""
+    uniform_samples = torch.rand(shape) * 0.99998 + 0.00001
+    return -torch.log(-torch.log(uniform_samples))
+
+
+def rand_gumbel_like(x):
+    g = rand_gumbel(x.size()).to(dtype=x.dtype, device=x.device)
+    return g
+
+
+def slice_segments(x, ids_str, segment_size=4):
+    ret = torch.zeros_like(x[:, :, :segment_size])
+    for i in range(x.size(0)):
+        idx_str = ids_str[i]
+        idx_end = idx_str + segment_size
+        ret[i] = x[i, :, idx_str:idx_end]
+    return ret
+
+
+def rand_slice_segments(x, x_lengths=None, segment_size=4):
+    b, d, t = x.size()
+    if x_lengths is None:
+        x_lengths = t
+    ids_str_max = x_lengths - segment_size + 1
+    ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
+    ret = slice_segments(x, ids_str, segment_size)
+    return ret, ids_str
+
+
+def get_timing_signal_1d(length, channels, min_timescale=1.0, max_timescale=1.0e4):
+    position = torch.arange(length, dtype=torch.float)
+    num_timescales = channels // 2
+    log_timescale_increment = math.log(float(max_timescale) / float(min_timescale)) / (
+        num_timescales - 1
+    )
+    inv_timescales = min_timescale * torch.exp(
+        torch.arange(num_timescales, dtype=torch.float) * -log_timescale_increment
+    )
+    scaled_time = position.unsqueeze(0) * inv_timescales.unsqueeze(1)
+    signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], 0)
+    signal = F.pad(signal, [0, 0, 0, channels % 2])
+    signal = signal.view(1, channels, length)
+    return signal
+
+
+def add_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4):
+    b, channels, length = x.size()
+    signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
+    return x + signal.to(dtype=x.dtype, device=x.device)
+
+
+def cat_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4, axis=1):
+    b, channels, length = x.size()
+    signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
+    return torch.cat([x, signal.to(dtype=x.dtype, device=x.device)], axis)
+
+
+def subsequent_mask(length):
+    mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
+    return mask
+
+
+@torch.jit.script
+def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
+    n_channels_int = n_channels[0]
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :n_channels_int, :])
+    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+    acts = t_act * s_act
+    return acts
+
+
+def convert_pad_shape(pad_shape):
+    layer = pad_shape[::-1]
+    pad_shape = [item for sublist in layer for item in sublist]
+    return pad_shape
+
+
+def shift_1d(x):
+    x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1]
+    return x
+
+
+def sequence_mask(length, max_length=None):
+    if max_length is None:
+        max_length = length.max()
+    x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+    return x.unsqueeze(0) < length.unsqueeze(1)
+
+
+def generate_path(duration, mask):
+    """
+    duration: [b, 1, t_x]
+    mask: [b, 1, t_y, t_x]
+    """
+
+    b, _, t_y, t_x = mask.shape
+    cum_duration = torch.cumsum(duration, -1)
+
+    cum_duration_flat = cum_duration.view(b * t_x)
+    path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype)
+    path = path.view(b, t_x, t_y)
+    path = path - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
+    path = path.unsqueeze(1).transpose(2, 3) * mask
+    return path
+
+
+def clip_grad_value_(parameters, clip_value, norm_type=2):
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = list(filter(lambda p: p.grad is not None, parameters))
+    norm_type = float(norm_type)
+    if clip_value is not None:
+        clip_value = float(clip_value)
+
+    total_norm = 0
+    for p in parameters:
+        param_norm = p.grad.data.norm(norm_type)
+        total_norm += param_norm.item() ** norm_type
+        if clip_value is not None:
+            p.grad.data.clamp_(min=-clip_value, max=clip_value)
+    total_norm = total_norm ** (1.0 / norm_type)
+    return total_norm

app/voiceclone/ov/mel_processing.py

@@ -0,0 +1,183 @@
+import torch
+import torch.utils.data
+from librosa.filters import mel as librosa_mel_fn
+
+MAX_WAV_VALUE = 32768.0
+
+
+def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
+    """
+    PARAMS
+    ------
+    C: compression factor
+    """
+    return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def dynamic_range_decompression_torch(x, C=1):
+    """
+    PARAMS
+    ------
+    C: compression factor used to compress
+    """
+    return torch.exp(x) / C
+
+
+def spectral_normalize_torch(magnitudes):
+    output = dynamic_range_compression_torch(magnitudes)
+    return output
+
+
+def spectral_de_normalize_torch(magnitudes):
+    output = dynamic_range_decompression_torch(magnitudes)
+    return output
+
+
+mel_basis = {}
+hann_window = {}
+
+
+def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False):
+    if torch.min(y) < -1.1:
+        print("min value is ", torch.min(y))
+    if torch.max(y) > 1.1:
+        print("max value is ", torch.max(y))
+
+    global hann_window
+    dtype_device = str(y.dtype) + "_" + str(y.device)
+    wnsize_dtype_device = str(win_size) + "_" + dtype_device
+    if wnsize_dtype_device not in hann_window:
+        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(
+            dtype=y.dtype, device=y.device
+        )
+
+    y = torch.nn.functional.pad(
+        y.unsqueeze(1),
+        (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),
+        mode="reflect",
+    )
+    y = y.squeeze(1)
+
+    spec = torch.stft(
+        y,
+        n_fft,
+        hop_length=hop_size,
+        win_length=win_size,
+        window=hann_window[wnsize_dtype_device],
+        center=center,
+        pad_mode="reflect",
+        normalized=False,
+        onesided=True,
+        return_complex=False,
+    )
+
+    spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
+    return spec
+
+
+def spectrogram_torch_conv(y, n_fft, sampling_rate, hop_size, win_size, center=False):
+    # if torch.min(y) < -1.:
+    #     print('min value is ', torch.min(y))
+    # if torch.max(y) > 1.:
+    #     print('max value is ', torch.max(y))
+
+    global hann_window
+    dtype_device = str(y.dtype) + '_' + str(y.device)
+    wnsize_dtype_device = str(win_size) + '_' + dtype_device
+    if wnsize_dtype_device not in hann_window:
+        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(dtype=y.dtype, device=y.device)
+
+    y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')
+    
+    # ******************** original ************************#
+    # y = y.squeeze(1)
+    # spec1 = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],
+    #                   center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=False)
+
+    # ******************** ConvSTFT ************************#
+    freq_cutoff = n_fft // 2 + 1
+    fourier_basis = torch.view_as_real(torch.fft.fft(torch.eye(n_fft)))
+    forward_basis = fourier_basis[:freq_cutoff].permute(2, 0, 1).reshape(-1, 1, fourier_basis.shape[1])
+    forward_basis = forward_basis * torch.as_tensor(librosa.util.pad_center(torch.hann_window(win_size), size=n_fft)).float()
+
+    import torch.nn.functional as F
+
+    # if center:
+    #     signal = F.pad(y[:, None, None, :], (n_fft // 2, n_fft // 2, 0, 0), mode = 'reflect').squeeze(1)
+    assert center is False
+
+    forward_transform_squared = F.conv1d(y, forward_basis.to(y.device), stride = hop_size)
+    spec2 = torch.stack([forward_transform_squared[:, :freq_cutoff, :], forward_transform_squared[:, freq_cutoff:, :]], dim = -1)
+
+
+    # ******************** Verification ************************#
+    spec1 = torch.stft(y.squeeze(1), n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],
+                      center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=False)
+    assert torch.allclose(spec1, spec2, atol=1e-4)
+
+    spec = torch.sqrt(spec2.pow(2).sum(-1) + 1e-6)
+    return spec
+
+
+def spec_to_mel_torch(spec, n_fft, num_mels, sampling_rate, fmin, fmax):
+    global mel_basis
+    dtype_device = str(spec.dtype) + "_" + str(spec.device)
+    fmax_dtype_device = str(fmax) + "_" + dtype_device
+    if fmax_dtype_device not in mel_basis:
+        mel = librosa_mel_fn(sampling_rate, n_fft, num_mels, fmin, fmax)
+        mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(
+            dtype=spec.dtype, device=spec.device
+        )
+    spec = torch.matmul(mel_basis[fmax_dtype_device], spec)
+    spec = spectral_normalize_torch(spec)
+    return spec
+
+
+def mel_spectrogram_torch(
+    y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False
+):
+    if torch.min(y) < -1.0:
+        print("min value is ", torch.min(y))
+    if torch.max(y) > 1.0:
+        print("max value is ", torch.max(y))
+
+    global mel_basis, hann_window
+    dtype_device = str(y.dtype) + "_" + str(y.device)
+    fmax_dtype_device = str(fmax) + "_" + dtype_device
+    wnsize_dtype_device = str(win_size) + "_" + dtype_device
+    if fmax_dtype_device not in mel_basis:
+        mel = librosa_mel_fn(sampling_rate, n_fft, num_mels, fmin, fmax)
+        mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(
+            dtype=y.dtype, device=y.device
+        )
+    if wnsize_dtype_device not in hann_window:
+        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(
+            dtype=y.dtype, device=y.device
+        )
+
+    y = torch.nn.functional.pad(
+        y.unsqueeze(1),
+        (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),
+        mode="reflect",
+    )
+    y = y.squeeze(1)
+
+    spec = torch.stft(
+        y,
+        n_fft,
+        hop_length=hop_size,
+        win_length=win_size,
+        window=hann_window[wnsize_dtype_device],
+        center=center,
+        pad_mode="reflect",
+        normalized=False,
+        onesided=True,
+        return_complex=False,
+    )
+
+    spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
+
+    spec = torch.matmul(mel_basis[fmax_dtype_device], spec)
+    spec = spectral_normalize_torch(spec)
+
+    return spec

app/voiceclone/ov/models.py

@@ -0,0 +1,499 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from . import commons
+from . import modules
+from . import attentions
+
+from torch.nn import Conv1d, ConvTranspose1d, Conv2d
+from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
+
+from .commons import init_weights, get_padding
+
+
+class TextEncoder(nn.Module):
+	def __init__(self,
+			n_vocab,
+			out_channels,
+			hidden_channels,
+			filter_channels,
+			n_heads,
+			n_layers,
+			kernel_size,
+			p_dropout):
+		super().__init__()
+		self.n_vocab = n_vocab
+		self.out_channels = out_channels
+		self.hidden_channels = hidden_channels
+		self.filter_channels = filter_channels
+		self.n_heads = n_heads
+		self.n_layers = n_layers
+		self.kernel_size = kernel_size
+		self.p_dropout = p_dropout
+
+		self.emb = nn.Embedding(n_vocab, hidden_channels)
+		nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5)
+
+		self.encoder = attentions.Encoder(
+			hidden_channels,
+			filter_channels,
+			n_heads,
+			n_layers,
+			kernel_size,
+			p_dropout)
+		self.proj= nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+	def forward(self, x, x_lengths):
+		x = self.emb(x) * math.sqrt(self.hidden_channels) # [b, t, h]
+		x = torch.transpose(x, 1, -1) # [b, h, t]
+		x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)
+
+		x = self.encoder(x * x_mask, x_mask)
+		stats = self.proj(x) * x_mask
+
+		m, logs = torch.split(stats, self.out_channels, dim=1)
+		return x, m, logs, x_mask
+     
+
+class DurationPredictor(nn.Module):
+    def __init__(
+        self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.gin_channels = gin_channels
+
+        self.drop = nn.Dropout(p_dropout)
+        self.conv_1 = nn.Conv1d(
+            in_channels, filter_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.norm_1 = modules.LayerNorm(filter_channels)
+        self.conv_2 = nn.Conv1d(
+            filter_channels, filter_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.norm_2 = modules.LayerNorm(filter_channels)
+        self.proj = nn.Conv1d(filter_channels, 1, 1)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, in_channels, 1)
+
+    def forward(self, x, x_mask, g=None):
+        x = torch.detach(x)
+        if g is not None:
+            g = torch.detach(g)
+            x = x + self.cond(g)
+        x = self.conv_1(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_1(x)
+        x = self.drop(x)
+        x = self.conv_2(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_2(x)
+        x = self.drop(x)
+        x = self.proj(x * x_mask)
+        return x * x_mask
+               
+class StochasticDurationPredictor(nn.Module):
+	def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, n_flows=4, gin_channels=0):
+		super().__init__()
+		filter_channels = in_channels # it needs to be removed from future version.
+		self.in_channels = in_channels
+		self.filter_channels = filter_channels
+		self.kernel_size = kernel_size
+		self.p_dropout = p_dropout
+		self.n_flows = n_flows
+		self.gin_channels = gin_channels
+
+		self.log_flow = modules.Log()
+		self.flows = nn.ModuleList()
+		self.flows.append(modules.ElementwiseAffine(2))
+		for i in range(n_flows):
+			self.flows.append(modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3))
+			self.flows.append(modules.Flip())
+
+		self.post_pre = nn.Conv1d(1, filter_channels, 1)
+		self.post_proj = nn.Conv1d(filter_channels, filter_channels, 1)
+		self.post_convs = modules.DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)
+		self.post_flows = nn.ModuleList()
+		self.post_flows.append(modules.ElementwiseAffine(2))
+		for i in range(4):
+			self.post_flows.append(modules.ConvFlow(2, filter_channels, kernel_size, n_layers=3))
+			self.post_flows.append(modules.Flip())
+
+		self.pre = nn.Conv1d(in_channels, filter_channels, 1)
+		self.proj = nn.Conv1d(filter_channels, filter_channels, 1)
+		self.convs = modules.DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)
+		if gin_channels != 0:
+			self.cond = nn.Conv1d(gin_channels, filter_channels, 1)
+
+	def forward(self, x, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):
+		x = torch.detach(x)
+		x = self.pre(x)
+		if g is not None:
+			g = torch.detach(g)
+			x = x + self.cond(g)
+		x = self.convs(x, x_mask)
+		x = self.proj(x) * x_mask
+
+		if not reverse:
+			flows = self.flows
+			assert w is not None
+
+			logdet_tot_q = 0
+			h_w = self.post_pre(w)
+			h_w = self.post_convs(h_w, x_mask)
+			h_w = self.post_proj(h_w) * x_mask
+			e_q = torch.randn(w.size(0), 2, w.size(2)).to(device=x.device, dtype=x.dtype) * x_mask
+			z_q = e_q
+			for flow in self.post_flows:
+				z_q, logdet_q = flow(z_q, x_mask, g=(x + h_w))
+				logdet_tot_q += logdet_q
+			z_u, z1 = torch.split(z_q, [1, 1], 1)
+			u = torch.sigmoid(z_u) * x_mask
+			z0 = (w - u) * x_mask
+			logdet_tot_q += torch.sum((F.logsigmoid(z_u) + F.logsigmoid(-z_u)) * x_mask, [1,2])
+			logq = torch.sum(-0.5 * (math.log(2*math.pi) + (e_q**2)) * x_mask, [1,2]) - logdet_tot_q
+
+			logdet_tot = 0
+			z0, logdet = self.log_flow(z0, x_mask)
+			logdet_tot += logdet
+			z = torch.cat([z0, z1], 1)
+			for flow in flows:
+				z, logdet = flow(z, x_mask, g=x, reverse=reverse)
+				logdet_tot = logdet_tot + logdet
+			nll = torch.sum(0.5 * (math.log(2*math.pi) + (z**2)) * x_mask, [1,2]) - logdet_tot
+			return nll + logq # [b]
+		else:
+			flows = list(reversed(self.flows))
+			flows = flows[:-2] + [flows[-1]] # remove a useless vflow
+			z = torch.randn(x.size(0), 2, x.size(2)).to(device=x.device, dtype=x.dtype) * noise_scale
+			for flow in flows:
+				z = flow(z, x_mask, g=x, reverse=reverse)
+			z0, z1 = torch.split(z, [1, 1], 1)
+			logw = z0
+			return logw
+
+class PosteriorEncoder(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        gin_channels=0,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+
+        self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
+        self.enc = modules.WN(
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            gin_channels=gin_channels,
+        )
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(self, x, x_lengths, g=None, tau=1.0):
+        x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(
+            x.dtype
+        )
+        x = self.pre(x) * x_mask
+        x = self.enc(x, x_mask, g=g)
+        stats = self.proj(x) * x_mask
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        z = (m + torch.randn_like(m) * tau * torch.exp(logs)) * x_mask
+        return z, m, logs, x_mask
+
+
+class Generator(torch.nn.Module):
+    def __init__(
+        self,
+        initial_channel,
+        resblock,
+        resblock_kernel_sizes,
+        resblock_dilation_sizes,
+        upsample_rates,
+        upsample_initial_channel,
+        upsample_kernel_sizes,
+        gin_channels=0,
+    ):
+        super(Generator, self).__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.conv_pre = Conv1d(
+            initial_channel, upsample_initial_channel, 7, 1, padding=3
+        )
+        resblock = modules.ResBlock1 if resblock == "1" else modules.ResBlock2
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(
+                weight_norm(
+                    ConvTranspose1d(
+                        upsample_initial_channel // (2**i),
+                        upsample_initial_channel // (2 ** (i + 1)),
+                        k,
+                        u,
+                        padding=(k - u) // 2,
+                    )
+                )
+            )
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel // (2 ** (i + 1))
+            for j, (k, d) in enumerate(
+                zip(resblock_kernel_sizes, resblock_dilation_sizes)
+            ):
+                self.resblocks.append(resblock(ch, k, d))
+
+        self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
+        self.ups.apply(init_weights)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
+
+    def forward(self, x, g=None):
+        x = self.conv_pre(x)
+        if g is not None:
+            x = x + self.cond(g)
+
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            x = self.ups[i](x)
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+
+        return x
+
+    def remove_weight_norm(self):
+        print("Removing weight norm...")
+        for layer in self.ups:
+            remove_weight_norm(layer)
+        for layer in self.resblocks:
+            layer.remove_weight_norm()
+
+
+class ReferenceEncoder(nn.Module):
+    """
+    inputs --- [N, Ty/r, n_mels*r]  mels
+    outputs --- [N, ref_enc_gru_size]
+    """
+
+    def __init__(self, spec_channels, gin_channels=0, layernorm=True):
+        super().__init__()
+        self.spec_channels = spec_channels
+        ref_enc_filters = [32, 32, 64, 64, 128, 128]
+        K = len(ref_enc_filters)
+        filters = [1] + ref_enc_filters
+        convs = [
+            weight_norm(
+                nn.Conv2d(
+                    in_channels=filters[i],
+                    out_channels=filters[i + 1],
+                    kernel_size=(3, 3),
+                    stride=(2, 2),
+                    padding=(1, 1),
+                )
+            )
+            for i in range(K)
+        ]
+        self.convs = nn.ModuleList(convs)
+
+        out_channels = self.calculate_channels(spec_channels, 3, 2, 1, K)
+        self.gru = nn.GRU(
+            input_size=ref_enc_filters[-1] * out_channels,
+            hidden_size=256 // 2,
+            batch_first=True,
+        )
+        self.proj = nn.Linear(128, gin_channels)
+        if layernorm:
+            self.layernorm = nn.LayerNorm(self.spec_channels)
+        else:
+            self.layernorm = None
+
+    def forward(self, inputs, mask=None):
+        N = inputs.size(0)
+
+        out = inputs.view(N, 1, -1, self.spec_channels)  # [N, 1, Ty, n_freqs]
+        if self.layernorm is not None:
+            out = self.layernorm(out)
+
+        for conv in self.convs:
+            out = conv(out)
+            # out = wn(out)
+            out = F.relu(out)  # [N, 128, Ty//2^K, n_mels//2^K]
+
+        out = out.transpose(1, 2)  # [N, Ty//2^K, 128, n_mels//2^K]
+        T = out.size(1)
+        N = out.size(0)
+        out = out.contiguous().view(N, T, -1)  # [N, Ty//2^K, 128*n_mels//2^K]
+
+        self.gru.flatten_parameters()
+        memory, out = self.gru(out)  # out --- [1, N, 128]
+
+        return self.proj(out.squeeze(0))
+
+    def calculate_channels(self, L, kernel_size, stride, pad, n_convs):
+        for i in range(n_convs):
+            L = (L - kernel_size + 2 * pad) // stride + 1
+        return L
+
+
+class ResidualCouplingBlock(nn.Module):
+    def __init__(self,
+            channels,
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            n_flows=4,
+            gin_channels=0):
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.flows = nn.ModuleList()
+        for i in range(n_flows):
+            self.flows.append(modules.ResidualCouplingLayer(channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels, mean_only=True))
+            self.flows.append(modules.Flip())
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        if not reverse:
+            for flow in self.flows:
+                x, _ = flow(x, x_mask, g=g, reverse=reverse)
+        else:
+            for flow in reversed(self.flows):
+                x = flow(x, x_mask, g=g, reverse=reverse)
+        return x
+
+class SynthesizerTrn(nn.Module):
+    """
+    Synthesizer for Training
+    """
+
+    def __init__(
+        self,
+        n_vocab,
+        spec_channels,
+        inter_channels,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size,
+        p_dropout,
+        resblock,
+        resblock_kernel_sizes,
+        resblock_dilation_sizes,
+        upsample_rates,
+        upsample_initial_channel,
+        upsample_kernel_sizes,
+        n_speakers=256,
+        gin_channels=256,
+        zero_g=False,
+        **kwargs
+    ):
+        super().__init__()
+
+        self.dec = Generator(
+            inter_channels,
+            resblock,
+            resblock_kernel_sizes,
+            resblock_dilation_sizes,
+            upsample_rates,
+            upsample_initial_channel,
+            upsample_kernel_sizes,
+            gin_channels=gin_channels,
+        )
+        self.enc_q = PosteriorEncoder(
+            spec_channels,
+            inter_channels,
+            hidden_channels,
+            5,
+            1,
+            16,
+            gin_channels=gin_channels,
+        )
+
+        self.flow = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, 4, gin_channels=gin_channels)
+
+        self.n_speakers = n_speakers
+        if n_speakers == 0:
+            self.ref_enc = ReferenceEncoder(spec_channels, gin_channels)
+        else:
+            self.enc_p = TextEncoder(n_vocab,
+                inter_channels,
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers,
+                kernel_size,
+                p_dropout)
+            self.sdp = StochasticDurationPredictor(hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels)
+            self.dp = DurationPredictor(hidden_channels, 256, 3, 0.5, gin_channels=gin_channels)
+            self.emb_g = nn.Embedding(n_speakers, gin_channels)
+        self.zero_g = zero_g
+
+    def infer(self, x, x_lengths, sid=None, noise_scale=1, length_scale=1, noise_scale_w=1., sdp_ratio=0.2, max_len=None):
+        x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths)
+        if self.n_speakers > 0:
+            g = self.emb_g(sid).unsqueeze(-1) # [b, h, 1]
+        else:
+            g = None
+
+        logw = self.sdp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w) * sdp_ratio \
+            + self.dp(x, x_mask, g=g) * (1 - sdp_ratio)
+
+        w = torch.exp(logw) * x_mask * length_scale
+        w_ceil = torch.ceil(w)
+        y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
+        y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, None), 1).to(x_mask.dtype)
+        attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
+        attn = commons.generate_path(w_ceil, attn_mask)
+
+        m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(1, 2) # [b, t', t], [b, t, d] -> [b, d, t']
+        logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(1, 2) # [b, t', t], [b, t, d] -> [b, d, t']
+
+        z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
+        z = self.flow(z_p, y_mask, g=g, reverse=True)
+        o = self.dec((z * y_mask)[:,:,:max_len], g=g)
+        return o, attn, y_mask, (z, z_p, m_p, logs_p)
+
+    def voice_conversion(self, y, y_lengths, sid_src, sid_tgt, tau=1.0):
+        g_src = sid_src
+        g_tgt = sid_tgt
+        z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g_src if not self.zero_g else torch.zeros_like(g_src), tau=tau)
+        z_p = self.flow(z, y_mask, g=g_src)
+        z_hat = self.flow(z_p, y_mask, g=g_tgt, reverse=True)
+        o_hat = self.dec(z_hat * y_mask, g=g_tgt if not self.zero_g else torch.zeros_like(g_tgt))
+        return o_hat, y_mask, (z, z_p, z_hat)

app/voiceclone/ov/modules.py

@@ -0,0 +1,598 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from torch.nn import Conv1d
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+from . import commons
+from .commons import init_weights, get_padding
+from .transforms import piecewise_rational_quadratic_transform
+from .attentions import Encoder
+
+LRELU_SLOPE = 0.1
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        x = x.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
+
+
+class ConvReluNorm(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        hidden_channels,
+        out_channels,
+        kernel_size,
+        n_layers,
+        p_dropout,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.p_dropout = p_dropout
+        assert n_layers > 1, "Number of layers should be larger than 0."
+
+        self.conv_layers = nn.ModuleList()
+        self.norm_layers = nn.ModuleList()
+        self.conv_layers.append(
+            nn.Conv1d(
+                in_channels, hidden_channels, kernel_size, padding=kernel_size // 2
+            )
+        )
+        self.norm_layers.append(LayerNorm(hidden_channels))
+        self.relu_drop = nn.Sequential(nn.ReLU(), nn.Dropout(p_dropout))
+        for _ in range(n_layers - 1):
+            self.conv_layers.append(
+                nn.Conv1d(
+                    hidden_channels,
+                    hidden_channels,
+                    kernel_size,
+                    padding=kernel_size // 2,
+                )
+            )
+            self.norm_layers.append(LayerNorm(hidden_channels))
+        self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
+        self.proj.weight.data.zero_()
+        self.proj.bias.data.zero_()
+
+    def forward(self, x, x_mask):
+        x_org = x
+        for i in range(self.n_layers):
+            x = self.conv_layers[i](x * x_mask)
+            x = self.norm_layers[i](x)
+            x = self.relu_drop(x)
+        x = x_org + self.proj(x)
+        return x * x_mask
+
+
+class DDSConv(nn.Module):
+    """
+    Dilated and Depth-Separable Convolution
+    """
+
+    def __init__(self, channels, kernel_size, n_layers, p_dropout=0.0):
+        super().__init__()
+        self.channels = channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.p_dropout = p_dropout
+
+        self.drop = nn.Dropout(p_dropout)
+        self.convs_sep = nn.ModuleList()
+        self.convs_1x1 = nn.ModuleList()
+        self.norms_1 = nn.ModuleList()
+        self.norms_2 = nn.ModuleList()
+        for i in range(n_layers):
+            dilation = kernel_size**i
+            padding = (kernel_size * dilation - dilation) // 2
+            self.convs_sep.append(
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    groups=channels,
+                    dilation=dilation,
+                    padding=padding,
+                )
+            )
+            self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
+            self.norms_1.append(LayerNorm(channels))
+            self.norms_2.append(LayerNorm(channels))
+
+    def forward(self, x, x_mask, g=None):
+        if g is not None:
+            x = x + g
+        for i in range(self.n_layers):
+            y = self.convs_sep[i](x * x_mask)
+            y = self.norms_1[i](y)
+            y = F.gelu(y)
+            y = self.convs_1x1[i](y)
+            y = self.norms_2[i](y)
+            y = F.gelu(y)
+            y = self.drop(y)
+            x = x + y
+        return x * x_mask
+
+
+class WN(torch.nn.Module):
+    def __init__(
+        self,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        gin_channels=0,
+        p_dropout=0,
+    ):
+        super(WN, self).__init__()
+        assert kernel_size % 2 == 1
+        self.hidden_channels = hidden_channels
+        self.kernel_size = (kernel_size,)
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+        self.p_dropout = p_dropout
+
+        self.in_layers = torch.nn.ModuleList()
+        self.res_skip_layers = torch.nn.ModuleList()
+        self.drop = nn.Dropout(p_dropout)
+
+        if gin_channels != 0:
+            cond_layer = torch.nn.Conv1d(
+                gin_channels, 2 * hidden_channels * n_layers, 1
+            )
+            self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name="weight")
+
+        for i in range(n_layers):
+            dilation = dilation_rate**i
+            padding = int((kernel_size * dilation - dilation) / 2)
+            in_layer = torch.nn.Conv1d(
+                hidden_channels,
+                2 * hidden_channels,
+                kernel_size,
+                dilation=dilation,
+                padding=padding,
+            )
+            in_layer = torch.nn.utils.weight_norm(in_layer, name="weight")
+            self.in_layers.append(in_layer)
+
+            # last one is not necessary
+            if i < n_layers - 1:
+                res_skip_channels = 2 * hidden_channels
+            else:
+                res_skip_channels = hidden_channels
+
+            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
+            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name="weight")
+            self.res_skip_layers.append(res_skip_layer)
+
+    def forward(self, x, x_mask, g=None, **kwargs):
+        output = torch.zeros_like(x)
+        n_channels_tensor = torch.IntTensor([self.hidden_channels])
+
+        if g is not None:
+            g = self.cond_layer(g)
+
+        for i in range(self.n_layers):
+            x_in = self.in_layers[i](x)
+            if g is not None:
+                cond_offset = i * 2 * self.hidden_channels
+                g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
+            else:
+                g_l = torch.zeros_like(x_in)
+
+            acts = commons.fused_add_tanh_sigmoid_multiply(x_in, g_l, n_channels_tensor)
+            acts = self.drop(acts)
+
+            res_skip_acts = self.res_skip_layers[i](acts)
+            if i < self.n_layers - 1:
+                res_acts = res_skip_acts[:, : self.hidden_channels, :]
+                x = (x + res_acts) * x_mask
+                output = output + res_skip_acts[:, self.hidden_channels :, :]
+            else:
+                output = output + res_skip_acts
+        return output * x_mask
+
+    def remove_weight_norm(self):
+        if self.gin_channels != 0:
+            torch.nn.utils.remove_weight_norm(self.cond_layer)
+        for l in self.in_layers:
+            torch.nn.utils.remove_weight_norm(l)
+        for l in self.res_skip_layers:
+            torch.nn.utils.remove_weight_norm(l)
+
+
+class ResBlock1(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super(ResBlock1, self).__init__()
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[2],
+                        padding=get_padding(kernel_size, dilation[2]),
+                    )
+                ),
+            ]
+        )
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+            ]
+        )
+        self.convs2.apply(init_weights)
+
+    def forward(self, x, x_mask=None):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c2(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+class ResBlock2(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
+        super(ResBlock2, self).__init__()
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+            ]
+        )
+        self.convs.apply(init_weights)
+
+    def forward(self, x, x_mask=None):
+        for c in self.convs:
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            remove_weight_norm(l)
+
+
+class Log(nn.Module):
+    def forward(self, x, x_mask, reverse=False, **kwargs):
+        if not reverse:
+            y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
+            logdet = torch.sum(-y, [1, 2])
+            return y, logdet
+        else:
+            x = torch.exp(x) * x_mask
+            return x
+
+
+class Flip(nn.Module):
+    def forward(self, x, *args, reverse=False, **kwargs):
+        x = torch.flip(x, [1])
+        if not reverse:
+            logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
+            return x, logdet
+        else:
+            return x
+
+
+class ElementwiseAffine(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.channels = channels
+        self.m = nn.Parameter(torch.zeros(channels, 1))
+        self.logs = nn.Parameter(torch.zeros(channels, 1))
+
+    def forward(self, x, x_mask, reverse=False, **kwargs):
+        if not reverse:
+            y = self.m + torch.exp(self.logs) * x
+            y = y * x_mask
+            logdet = torch.sum(self.logs * x_mask, [1, 2])
+            return y, logdet
+        else:
+            x = (x - self.m) * torch.exp(-self.logs) * x_mask
+            return x
+
+
+class ResidualCouplingLayer(nn.Module):
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        p_dropout=0,
+        gin_channels=0,
+        mean_only=False,
+    ):
+        assert channels % 2 == 0, "channels should be divisible by 2"
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.half_channels = channels // 2
+        self.mean_only = mean_only
+
+        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+        self.enc = WN(
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            p_dropout=p_dropout,
+            gin_channels=gin_channels,
+        )
+        self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+        self.post.weight.data.zero_()
+        self.post.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0) * x_mask
+        h = self.enc(h, x_mask, g=g)
+        stats = self.post(h) * x_mask
+        if not self.mean_only:
+            m, logs = torch.split(stats, [self.half_channels] * 2, 1)
+        else:
+            m = stats
+            logs = torch.zeros_like(m)
+
+        if not reverse:
+            x1 = m + x1 * torch.exp(logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            logdet = torch.sum(logs, [1, 2])
+            return x, logdet
+        else:
+            x1 = (x1 - m) * torch.exp(-logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            return x
+
+
+class ConvFlow(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        filter_channels,
+        kernel_size,
+        n_layers,
+        num_bins=10,
+        tail_bound=5.0,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.num_bins = num_bins
+        self.tail_bound = tail_bound
+        self.half_channels = in_channels // 2
+
+        self.pre = nn.Conv1d(self.half_channels, filter_channels, 1)
+        self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.0)
+        self.proj = nn.Conv1d(
+            filter_channels, self.half_channels * (num_bins * 3 - 1), 1
+        )
+        self.proj.weight.data.zero_()
+        self.proj.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0)
+        h = self.convs(h, x_mask, g=g)
+        h = self.proj(h) * x_mask
+
+        b, c, t = x0.shape
+        h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2)  # [b, cx?, t] -> [b, c, t, ?]
+
+        unnormalized_widths = h[..., : self.num_bins] / math.sqrt(self.filter_channels)
+        unnormalized_heights = h[..., self.num_bins : 2 * self.num_bins] / math.sqrt(
+            self.filter_channels
+        )
+        unnormalized_derivatives = h[..., 2 * self.num_bins :]
+
+        x1, logabsdet = piecewise_rational_quadratic_transform(
+            x1,
+            unnormalized_widths,
+            unnormalized_heights,
+            unnormalized_derivatives,
+            inverse=reverse,
+            tails="linear",
+            tail_bound=self.tail_bound,
+        )
+
+        x = torch.cat([x0, x1], 1) * x_mask
+        logdet = torch.sum(logabsdet * x_mask, [1, 2])
+        if not reverse:
+            return x, logdet
+        else:
+            return x
+
+
+class TransformerCouplingLayer(nn.Module):
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        kernel_size,
+        n_layers,
+        n_heads,
+        p_dropout=0,
+        filter_channels=0,
+        mean_only=False,
+        wn_sharing_parameter=None,
+        gin_channels=0,
+    ):
+        assert n_layers == 3, n_layers
+        assert channels % 2 == 0, "channels should be divisible by 2"
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.half_channels = channels // 2
+        self.mean_only = mean_only
+
+        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+        self.enc = (
+            Encoder(
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers,
+                kernel_size,
+                p_dropout,
+                isflow=True,
+                gin_channels=gin_channels,
+            )
+            if wn_sharing_parameter is None
+            else wn_sharing_parameter
+        )
+        self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+        self.post.weight.data.zero_()
+        self.post.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0) * x_mask
+        h = self.enc(h, x_mask, g=g)
+        stats = self.post(h) * x_mask
+        if not self.mean_only:
+            m, logs = torch.split(stats, [self.half_channels] * 2, 1)
+        else:
+            m = stats
+            logs = torch.zeros_like(m)
+
+        if not reverse:
+            x1 = m + x1 * torch.exp(logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            logdet = torch.sum(logs, [1, 2])
+            return x, logdet
+        else:
+            x1 = (x1 - m) * torch.exp(-logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            return x
+
+        x1, logabsdet = piecewise_rational_quadratic_transform(
+            x1,
+            unnormalized_widths,
+            unnormalized_heights,
+            unnormalized_derivatives,
+            inverse=reverse,
+            tails="linear",
+            tail_bound=self.tail_bound,
+        )
+
+        x = torch.cat([x0, x1], 1) * x_mask
+        logdet = torch.sum(logabsdet * x_mask, [1, 2])
+        if not reverse:
+            return x, logdet
+        else:
+            return x

app/voiceclone/ov/transforms.py

@@ -0,0 +1,209 @@
+import torch
+from torch.nn import functional as F
+
+import numpy as np
+
+
+DEFAULT_MIN_BIN_WIDTH = 1e-3
+DEFAULT_MIN_BIN_HEIGHT = 1e-3
+DEFAULT_MIN_DERIVATIVE = 1e-3
+
+
+def piecewise_rational_quadratic_transform(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    tails=None,
+    tail_bound=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    if tails is None:
+        spline_fn = rational_quadratic_spline
+        spline_kwargs = {}
+    else:
+        spline_fn = unconstrained_rational_quadratic_spline
+        spline_kwargs = {"tails": tails, "tail_bound": tail_bound}
+
+    outputs, logabsdet = spline_fn(
+        inputs=inputs,
+        unnormalized_widths=unnormalized_widths,
+        unnormalized_heights=unnormalized_heights,
+        unnormalized_derivatives=unnormalized_derivatives,
+        inverse=inverse,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+        **spline_kwargs
+    )
+    return outputs, logabsdet
+
+
+def searchsorted(bin_locations, inputs, eps=1e-6):
+    bin_locations[..., -1] += eps
+    return torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1
+
+
+def unconstrained_rational_quadratic_spline(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    tails="linear",
+    tail_bound=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
+    outside_interval_mask = ~inside_interval_mask
+
+    outputs = torch.zeros_like(inputs)
+    logabsdet = torch.zeros_like(inputs)
+
+    if tails == "linear":
+        unnormalized_derivatives = F.pad(unnormalized_derivatives, pad=(1, 1))
+        constant = np.log(np.exp(1 - min_derivative) - 1)
+        unnormalized_derivatives[..., 0] = constant
+        unnormalized_derivatives[..., -1] = constant
+
+        outputs[outside_interval_mask] = inputs[outside_interval_mask]
+        logabsdet[outside_interval_mask] = 0
+    else:
+        raise RuntimeError("{} tails are not implemented.".format(tails))
+
+    (
+        outputs[inside_interval_mask],
+        logabsdet[inside_interval_mask],
+    ) = rational_quadratic_spline(
+        inputs=inputs[inside_interval_mask],
+        unnormalized_widths=unnormalized_widths[inside_interval_mask, :],
+        unnormalized_heights=unnormalized_heights[inside_interval_mask, :],
+        unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :],
+        inverse=inverse,
+        left=-tail_bound,
+        right=tail_bound,
+        bottom=-tail_bound,
+        top=tail_bound,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+    )
+
+    return outputs, logabsdet
+
+
+def rational_quadratic_spline(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    left=0.0,
+    right=1.0,
+    bottom=0.0,
+    top=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    if torch.min(inputs) < left or torch.max(inputs) > right:
+        raise ValueError("Input to a transform is not within its domain")
+
+    num_bins = unnormalized_widths.shape[-1]
+
+    if min_bin_width * num_bins > 1.0:
+        raise ValueError("Minimal bin width too large for the number of bins")
+    if min_bin_height * num_bins > 1.0:
+        raise ValueError("Minimal bin height too large for the number of bins")
+
+    widths = F.softmax(unnormalized_widths, dim=-1)
+    widths = min_bin_width + (1 - min_bin_width * num_bins) * widths
+    cumwidths = torch.cumsum(widths, dim=-1)
+    cumwidths = F.pad(cumwidths, pad=(1, 0), mode="constant", value=0.0)
+    cumwidths = (right - left) * cumwidths + left
+    cumwidths[..., 0] = left
+    cumwidths[..., -1] = right
+    widths = cumwidths[..., 1:] - cumwidths[..., :-1]
+
+    derivatives = min_derivative + F.softplus(unnormalized_derivatives)
+
+    heights = F.softmax(unnormalized_heights, dim=-1)
+    heights = min_bin_height + (1 - min_bin_height * num_bins) * heights
+    cumheights = torch.cumsum(heights, dim=-1)
+    cumheights = F.pad(cumheights, pad=(1, 0), mode="constant", value=0.0)
+    cumheights = (top - bottom) * cumheights + bottom
+    cumheights[..., 0] = bottom
+    cumheights[..., -1] = top
+    heights = cumheights[..., 1:] - cumheights[..., :-1]
+
+    if inverse:
+        bin_idx = searchsorted(cumheights, inputs)[..., None]
+    else:
+        bin_idx = searchsorted(cumwidths, inputs)[..., None]
+
+    input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0]
+    input_bin_widths = widths.gather(-1, bin_idx)[..., 0]
+
+    input_cumheights = cumheights.gather(-1, bin_idx)[..., 0]
+    delta = heights / widths
+    input_delta = delta.gather(-1, bin_idx)[..., 0]
+
+    input_derivatives = derivatives.gather(-1, bin_idx)[..., 0]
+    input_derivatives_plus_one = derivatives[..., 1:].gather(-1, bin_idx)[..., 0]
+
+    input_heights = heights.gather(-1, bin_idx)[..., 0]
+
+    if inverse:
+        a = (inputs - input_cumheights) * (
+            input_derivatives + input_derivatives_plus_one - 2 * input_delta
+        ) + input_heights * (input_delta - input_derivatives)
+        b = input_heights * input_derivatives - (inputs - input_cumheights) * (
+            input_derivatives + input_derivatives_plus_one - 2 * input_delta
+        )
+        c = -input_delta * (inputs - input_cumheights)
+
+        discriminant = b.pow(2) - 4 * a * c
+        assert (discriminant >= 0).all()
+
+        root = (2 * c) / (-b - torch.sqrt(discriminant))
+        outputs = root * input_bin_widths + input_cumwidths
+
+        theta_one_minus_theta = root * (1 - root)
+        denominator = input_delta + (
+            (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+            * theta_one_minus_theta
+        )
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * root.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - root).pow(2)
+        )
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, -logabsdet
+    else:
+        theta = (inputs - input_cumwidths) / input_bin_widths
+        theta_one_minus_theta = theta * (1 - theta)
+
+        numerator = input_heights * (
+            input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta
+        )
+        denominator = input_delta + (
+            (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+            * theta_one_minus_theta
+        )
+        outputs = input_cumheights + numerator / denominator
+
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * theta.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - theta).pow(2)
+        )
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, logabsdet

app/voiceclone/ov/utils.py

@@ -0,0 +1,194 @@
+import re
+import json
+import numpy as np
+
+
+def get_hparams_from_file(config_path):
+    with open(config_path, "r", encoding="utf-8") as f:
+        data = f.read()
+    config = json.loads(data)
+
+    hparams = HParams(**config)
+    return hparams
+
+class HParams:
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            if type(v) == dict:
+                v = HParams(**v)
+            self[k] = v
+
+    def keys(self):
+        return self.__dict__.keys()
+
+    def items(self):
+        return self.__dict__.items()
+
+    def values(self):
+        return self.__dict__.values()
+
+    def __len__(self):
+        return len(self.__dict__)
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def __setitem__(self, key, value):
+        return setattr(self, key, value)
+
+    def __contains__(self, key):
+        return key in self.__dict__
+
+    def __repr__(self):
+        return self.__dict__.__repr__()
+
+
+def string_to_bits(string, pad_len=8):
+    # Convert each character to its ASCII value
+    ascii_values = [ord(char) for char in string]
+    
+    # Convert ASCII values to binary representation
+    binary_values = [bin(value)[2:].zfill(8) for value in ascii_values]
+    
+    # Convert binary strings to integer arrays
+    bit_arrays = [[int(bit) for bit in binary] for binary in binary_values]
+    
+    # Convert list of arrays to NumPy array
+    numpy_array = np.array(bit_arrays)
+    numpy_array_full = np.zeros((pad_len, 8), dtype=numpy_array.dtype)
+    numpy_array_full[:, 2] = 1
+    max_len = min(pad_len, len(numpy_array))
+    numpy_array_full[:max_len] = numpy_array[:max_len]
+    return numpy_array_full
+
+
+def bits_to_string(bits_array):
+    # Convert each row of the array to a binary string
+    binary_values = [''.join(str(bit) for bit in row) for row in bits_array]
+    
+    # Convert binary strings to ASCII values
+    ascii_values = [int(binary, 2) for binary in binary_values]
+    
+    # Convert ASCII values to characters
+    output_string = ''.join(chr(value) for value in ascii_values)
+    
+    return output_string
+
+
+def split_sentence(text, min_len=10, language_str='[EN]'):
+    if language_str in ['EN']:
+        sentences = split_sentences_latin(text, min_len=min_len)
+    else:
+        sentences = split_sentences_zh(text, min_len=min_len)
+    return sentences
+
+def split_sentences_latin(text, min_len=10):
+    """Split Long sentences into list of short ones
+
+    Args:
+        str: Input sentences.
+
+    Returns:
+        List[str]: list of output sentences.
+    """
+    # deal with dirty sentences
+    text = re.sub('[。！？；]', '.', text)
+    text = re.sub('[，]', ',', text)
+    text = re.sub('[“”]', '"', text)
+    text = re.sub('[‘’]', "'", text)
+    text = re.sub(r"[\<\>\(\)\[\]\"\«\»]+", "", text)
+    text = re.sub('[\n\t ]+', ' ', text)
+    text = re.sub('([,.!?;])', r'\1 $#!', text)
+    # split
+    sentences = [s.strip() for s in text.split('$#!')]
+    if len(sentences[-1]) == 0: del sentences[-1]
+
+    new_sentences = []
+    new_sent = []
+    count_len = 0
+    for ind, sent in enumerate(sentences):
+        # print(sent)
+        new_sent.append(sent)
+        count_len += len(sent.split(" "))
+        if count_len > min_len or ind == len(sentences) - 1:
+            count_len = 0
+            new_sentences.append(' '.join(new_sent))
+            new_sent = []
+    return merge_short_sentences_latin(new_sentences)
+
+
+def merge_short_sentences_latin(sens):
+    """Avoid short sentences by merging them with the following sentence.
+
+    Args:
+        List[str]: list of input sentences.
+
+    Returns:
+        List[str]: list of output sentences.
+    """
+    sens_out = []
+    for s in sens:
+        # If the previous sentence is too short, merge them with
+        # the current sentence.
+        if len(sens_out) > 0 and len(sens_out[-1].split(" ")) <= 2:
+            sens_out[-1] = sens_out[-1] + " " + s
+        else:
+            sens_out.append(s)
+    try:
+        if len(sens_out[-1].split(" ")) <= 2:
+            sens_out[-2] = sens_out[-2] + " " + sens_out[-1]
+            sens_out.pop(-1)
+    except:
+        pass
+    return sens_out
+
+def split_sentences_zh(text, min_len=10):
+    text = re.sub('[。！？；]', '.', text)
+    text = re.sub('[，]', ',', text)
+    # 将文本中的换行符、空格和制表符替换为空格
+    text = re.sub('[\n\t ]+', ' ', text)
+    # 在标点符号后添加一个空格
+    text = re.sub('([,.!?;])', r'\1 $#!', text)
+    # 分隔句子并去除前后空格
+    # sentences = [s.strip() for s in re.split('(。|！|？|；)', text)]
+    sentences = [s.strip() for s in text.split('$#!')]
+    if len(sentences[-1]) == 0: del sentences[-1]
+
+    new_sentences = []
+    new_sent = []
+    count_len = 0
+    for ind, sent in enumerate(sentences):
+        new_sent.append(sent)
+        count_len += len(sent)
+        if count_len > min_len or ind == len(sentences) - 1:
+            count_len = 0
+            new_sentences.append(' '.join(new_sent))
+            new_sent = []
+    return merge_short_sentences_zh(new_sentences)
+
+
+def merge_short_sentences_zh(sens):
+    # return sens
+    """Avoid short sentences by merging them with the following sentence.
+
+    Args:
+        List[str]: list of input sentences.
+
+    Returns:
+        List[str]: list of output sentences.
+    """
+    sens_out = []
+    for s in sens:
+        # If the previous sentense is too short, merge them with
+        # the current sentence.
+        if len(sens_out) > 0 and len(sens_out[-1]) <= 2:
+            sens_out[-1] = sens_out[-1] + " " + s
+        else:
+            sens_out.append(s)
+    try:
+        if len(sens_out[-1]) <= 2:
+            sens_out[-2] = sens_out[-2] + " " + sens_out[-1]
+            sens_out.pop(-1)
+    except:
+        pass
+    return sens_out

requirements.txt

@@ -11,5 +11,8 @@ numpy
 edge-tts>=7.0
 deep-translator>=1.11
 requests
+# voice cloning (OpenVoice V2 converter; torch CPU is installed in the Dockerfile)
+librosa>=0.10
+soundfile
 # optional, only if you set ANTHROPIC_API_KEY
 anthropic

train/finetune_nllb.py

@@ -0,0 +1,117 @@
+"""Fine-tune the translation model on what people taught SwitchLingo.
+
+Input: the feedback export (GET /api/feedback/export?token=...), which contains
+  - viewer corrections   {original, correction, src, tgt}
+  - taught sentences     {kind: "pair", en, text, lang, dialect}
+Output: a fine-tuned NLLB model, converted to CTranslate2 int8 and (optionally) pushed to the Hub.
+Point the site at it with the Space variable OXT_NLLB_MODEL=<your-repo>.
+
+Run on a GPU (Google Colab free T4 or Hugging Face Jobs). A few thousand sentences train in ~15 minutes.
+
+  pip install torch transformers datasets sentencepiece ctranslate2 huggingface_hub accelerate
+  python train/finetune_nllb.py feedback.jsonl --langs lo hmn --push MsThai/switchlingo-nllb-lo
+  python train/finetune_nllb.py feedback.jsonl --dry-run      # just show what would be trained
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import random
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from app.languages import LANGS  # noqa: E402
+
+BASE = "facebook/nllb-200-distilled-600M"
+
+
+def load_pairs(path: str, langs: set[str] | None, min_votes: int) -> list[tuple[str, str, str, str]]:
+    """(src_lang, tgt_lang, source_text, target_text). Corrections need min_votes distinct voters
+    unless trusted; taught pairs are used both directions."""
+    votes: dict[tuple, set] = {}
+    trusted: set[tuple] = set()
+    for line in open(path, encoding="utf-8"):
+        try:
+            r = json.loads(line)
+        except ValueError:
+            continue
+        if r.get("kind") == "pair":
+            if r.get("dialect"):
+                continue  # dialects get their own model later
+            items = [("en", r["lang"], r["en"], r["text"]), (r["lang"], "en", r["text"], r["en"])]
+        elif r.get("correction"):
+            items = [(r["src"], r["tgt"], r["original"], r["correction"])]
+        else:
+            continue
+        for it in items:
+            it = tuple(x.strip() for x in it)
+            if not all(it):
+                continue
+            votes.setdefault(it, set()).add(r.get("voter", "?"))
+            if r.get("trusted"):
+                trusted.add(it)
+    out = [k for k, v in votes.items() if k in trusted or len(v) >= min_votes]
+    if langs:
+        out = [p for p in out if p[0] in langs or p[1] in langs]
+    return [p for p in out if LANGS.get(p[0], {}).get("nllb") and LANGS.get(p[1], {}).get("nllb")]
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("export")
+    ap.add_argument("--langs", nargs="*", help="only pairs touching these codes, e.g. lo hmn")
+    ap.add_argument("--min-votes", type=int, default=2)
+    ap.add_argument("--epochs", type=float, default=3)
+    ap.add_argument("--lr", type=float, default=5e-5)
+    ap.add_argument("--out", default="nllb-switchlingo")
+    ap.add_argument("--push", help="Hub repo for the CTranslate2 model")
+    ap.add_argument("--dry-run", action="store_true")
+    a = ap.parse_args()
+
+    pairs = load_pairs(a.export, set(a.langs) if a.langs else None, a.min_votes)
+    by = {}
+    for p in pairs:
+        by[f"{p[0]}→{p[1]}"] = by.get(f"{p[0]}→{p[1]}", 0) + 1
+    print(f"{len(pairs)} training pairs:", by)
+    if a.dry_run or not pairs:
+        for p in pairs[:5]:
+            print("  ", p)
+        return
+
+    import torch
+    from datasets import Dataset
+    from transformers import (AutoModelForSeq2SeqLM, AutoTokenizer, DataCollatorForSeq2Seq,
+                              Seq2SeqTrainer, Seq2SeqTrainingArguments)
+    tok = AutoTokenizer.from_pretrained(BASE)
+    model = AutoModelForSeq2SeqLM.from_pretrained(BASE)
+    random.shuffle(pairs)
+
+    def enc(ex):
+        tok.src_lang, tok.tgt_lang = LANGS[ex["s"]]["nllb"], LANGS[ex["t"]]["nllb"]
+        return tok(ex["x"], text_target=ex["y"], truncation=True, max_length=200)
+
+    ds = Dataset.from_list([{"s": s, "t": t, "x": x, "y": y} for s, t, x, y in pairs]).map(
+        enc, remove_columns=["s", "t", "x", "y"])
+    args = Seq2SeqTrainingArguments(a.out, num_train_epochs=a.epochs, learning_rate=a.lr,
+                                    per_device_train_batch_size=8, gradient_accumulation_steps=2,
+                                    warmup_ratio=0.1, weight_decay=0.01, save_strategy="no",
+                                    fp16=torch.cuda.is_available(), logging_steps=20, report_to=[])
+    Seq2SeqTrainer(model=model, args=args, train_dataset=ds,
+                   data_collator=DataCollatorForSeq2Seq(tok, model)).train()
+    model.save_pretrained(a.out)
+    tok.save_pretrained(a.out)
+
+    import ctranslate2.converters
+    ct2 = a.out + "-ct2-int8"
+    ctranslate2.converters.TransformersConverter(a.out).convert(ct2, quantization="int8", force=True)
+    print("CTranslate2 model:", ct2)
+    if a.push:
+        from huggingface_hub import HfApi
+        HfApi().create_repo(a.push, exist_ok=True, private=True)
+        HfApi().upload_folder(folder_path=ct2, repo_id=a.push)
+        print(f"Pushed. Set the Space variable OXT_NLLB_MODEL={a.push}")
+
+
+if __name__ == "__main__":
+    main()