Replace ZeroGPU demo with Gradio studio and H200 endpoint

.dockerignore

@@ -0,0 +1,8 @@
+.git
+.venv
+__pycache__
+.pytest_cache
+**/__pycache__
+**/*.pyc
+endpoint/comparisons
+endpoint/results

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+.pytest_cache/
+.venv/
+*.pyc

Dockerfile

@@ -0,0 +1,23 @@
+# The Space hosts the Gradio UI on CPU. The pinned H200 endpoint serves Drama 3.
+# Dockerfile.engine retains the GPU image recipe used by that endpoint.
+FROM python:3.11-slim-bookworm
+RUN apt-get update && apt-get install -y --no-install-recommends libgomp1 \
+    && rm -rf /var/lib/apt/lists/* \
+    && useradd -m -u 1000 user
+WORKDIR /home/user/app
+COPY requirements-ui.txt requirements-preview.txt /tmp/
+RUN pip install --no-cache-dir -r /tmp/requirements-ui.txt
+ENV HOME=/home/user \
+    USER=user \
+    PYTHONUNBUFFERED=1 \
+    GRADIO_ANALYTICS_ENABLED=False \
+    WHISPER_DIR=/home/user/.cache/whisper \
+    DEMO_UNLIMITED_ORGS=fishaudio \
+    USAGE_BUCKET_MOUNT_PATH=/data \
+    DRAMA_ENDPOINT_URL=https://6acfbd40fda97227704d0f72.endpoints.huggingface.cloud
+USER user
+# Cache the same Whisper turbo model for CPU transcription, independently of UI edits.
+RUN python -c "from faster_whisper import download_model; download_model('turbo', cache_dir='/home/user/.cache/whisper')"
+COPY --chown=user:user . .
+EXPOSE 7860
+CMD ["python", "server.py"]

Dockerfile.engine

@@ -0,0 +1,65 @@
+# Fish Audio Drama 3 on SGLang Omni, as a Hugging Face Docker Space on one GPU
+# (RTX Pro 6000 or H200).
+#
+# The base image is the one SGLang Omni's CI runs main on: CUDA, SGLang, and most
+# Python dependencies, but not sglang-omni itself, which is installed at a pinned commit.
+FROM hongccc/sglang-omni@sha256:ebe4239e29a764ee3a2806385c061c5fd438a26f01458e503d3822dcba5790df
+
+ARG SGLANG_OMNI_COMMIT=22d4ace88f9547777b819a69d7691204bb9ce847
+
+# The image moves slower than SGLang Omni (its SGLang predates the pinned commit's
+# sglang==0.5.21). Install the commit's pins the image does not satisfy, as SGLang
+# Omni's CI does before testing on this image (.github/scripts/prepare_omni_venv.sh).
+COPY scripts/omni_missing_deps.py /tmp/omni_missing_deps.py
+RUN python3 -c "import sys, urllib.request; urllib.request.urlretrieve(sys.argv[1], '/tmp/omni-pyproject.toml')" \
+        "https://raw.githubusercontent.com/sgl-project/sglang-omni/${SGLANG_OMNI_COMMIT}/pyproject.toml" \
+    && python3 /tmp/omni_missing_deps.py /tmp/omni-pyproject.toml > /tmp/omni-missing.txt \
+    && echo "Pins the image does not satisfy:" && cat /tmp/omni-missing.txt \
+    && if [ -s /tmp/omni-missing.txt ]; then xargs -d '\n' -n 1 python3 -m pip install --no-cache-dir < /tmp/omni-missing.txt; fi \
+    && python3 /tmp/omni_missing_deps.py --overrides /tmp/omni-pyproject.toml \
+        | xargs -r -d '\n' python3 -m pip install --no-cache-dir --no-deps \
+    && python3 /tmp/omni_missing_deps.py --report /tmp/omni-pyproject.toml
+
+RUN python3 -m pip install --no-cache-dir --no-deps \
+        "sglang-omni @ https://github.com/sgl-project/sglang-omni/archive/${SGLANG_OMNI_COMMIT}.tar.gz" \
+        python-multipart \
+    && python3 -m pip install --no-cache-dir authlib itsdangerous "gradio[oauth]==6.28.0" \
+    && python3 -c "import sglang_omni, whisper, av, librosa, httpx, fastapi, huggingface_hub; from huggingface_hub import attach_huggingface_oauth, parse_huggingface_oauth; print('sglang-omni', sglang_omni.__version__, 'huggingface_hub', huggingface_hub.__version__)"
+
+# SGLang Omni declares each stage's factory in the model's config class and refuses
+# to override it from the command line. Rebind the class name the Drama 3
+# preprocessing factory builds to the Space's subclass, which adds described
+# characters and directing history (drama3_space_ext.py, on PYTHONPATH at runtime).
+RUN stages=$(python3 -c "import os, sglang_omni; print(os.path.join(os.path.dirname(sglang_omni.__file__), 'models', 'fishaudio_drama3', 'stages.py'))") \
+    && grep -q "^def create_preprocessing_executor" "$stages" \
+    && grep -q "preprocessor = Drama3SpeechPreprocessor(" "$stages" \
+    && printf '\n\n# Drama 3 Space: described characters and directing history.\nfrom drama3_space_ext import Drama3SpacePreprocessor as Drama3SpeechPreprocessor  # noqa: E402\n' >> "$stages" \
+    && tail -3 "$stages"
+
+# Spaces run the container as uid 1000. Give it a writable home that holds the
+# FlashInfer JIT cache the image built for Hopper, and stop the base entrypoint from
+# cloning SGLang Omni's main branch at start.
+RUN (id -u 1000 >/dev/null 2>&1 || useradd -m -u 1000 -s /bin/bash user) \
+    && mkdir -p /home/user/app /home/user/.cache \
+    && if [ -d /root/.cache/flashinfer ]; then cp -a /root/.cache/flashinfer /home/user/.cache/; fi \
+    && chown -R 1000:1000 /home/user
+ENV HOME=/home/user \
+    USER=user \
+    FLASHINFER_WORKSPACE_BASE=/home/user \
+    SGLANG_OMNI_AUTO_CLONE=0 \
+    PYTHONUNBUFFERED=1 \
+    PYTHONPATH=/home/user/app \
+    WHISPER_DIR=/home/user/.cache/whisper \
+    DEMO_UNLIMITED_ORGS=fishaudio \
+    USAGE_BUCKET_MOUNT_PATH=/data
+
+USER 1000
+WORKDIR /home/user/app
+
+# Whisper large-v3-turbo transcribes reference clips (about 1.5 GB).
+RUN python3 -c "import whisper; whisper._download(whisper._MODELS['turbo'], '/home/user/.cache/whisper', False)"
+
+COPY --chown=1000:1000 . .
+
+EXPOSE 7860
+CMD ["python3", "server.py"]

README.md

@@ -1,107 +1,245 @@
 ---
 title: Drama 3
-emoji: 👀
+emoji: 🎭
 colorFrom: purple
 colorTo: red
-sdk: gradio
-sdk_version: 6.30.0
-python_version: '3.12'
-app_file: app.py
+sdk: docker
+app_port: 7860
+hf_oauth: true
 pinned: false
-preload_from_hub:
-  - fishaudio/drama-3
-startup_duration_timeout: 1h
+license: apache-2.0
+short_description: Direct Fish Audio Drama 3, streaming from one GPU
+startup_duration_timeout: 30m
 models:
   - fishaudio/drama-3
-license: apache-2.0
-short_description: Fish Audio Drama 3 Work Space!
+  - fishaudio/s2-pro
 tags:
-  - ml-intern
   - text-to-speech
   - fish-audio
-  - zerogpu
+  - sglang-omni
+  - voice-cloning
+  - voice-design
+  - streaming
 ---
 
-[![Built with ML Intern](https://huggingface.co/datasets/huggingface/badges/resolve/main/built-with-ml-intern-sm.svg)](https://hf.co/chat?mode=ml-intern)
-
-# Fish Audio Drama 3 — single-speaker voice cloning
-
-A Gradio app that runs
-**[fishaudio/drama-3](https://huggingface.co/fishaudio/drama-3)** (35B MoE, 3B
-active per token) **in-process on a ZeroGPU `large` slice** (RTX Pro 6000, 48 GB)
-via one `@spaces.GPU` function. Record or upload a voice sample, type what that
-sample says, and generate a new line in that voice.
-
-> **Hardware:** this Space must run on **ZeroGPU** hardware (Settings → Hardware → ZeroGPU).
-> On CPU hardware the page opens, but Generate asks you to flip the hardware.
-
-Inference uses the official PyTorch path of the model, vendored under `fish_speech/`
-(adapted from [hugging-apps/fish-audio-drama-3](https://huggingface.co/spaces/hugging-apps/fish-audio-drama-3)
-via [akhaliq/drama-3-tts](https://huggingface.co/spaces/akhaliq/drama-3-tts)):
-the 32.2B FP8 MoE experts stay in FP8 — dequantised they would need ~67 GB and
-would not fit the slice. The model loads at module scope, as the ZeroGPU docs
-recommend: built on the meta device, materialised on CPU with FP8 experts, then
-moved to `cuda`, so startup costs no GPU time. `spaces` is imported before CUDA
-is initialized.
-
-## How to use
-
-The page has five controls:
-
-1. **Reference audio** — upload a clip or record with the microphone. The clip
-   must be 0.5–30 seconds of speech. A longer clip is rejected; it is not
-   silently trimmed, because the transcript has to match the whole recording.
-2. **Reference transcript** — type exactly what is spoken in the reference
-   audio. This text only conditions the voice. It is not the line that gets
-   synthesized.
-3. **Text to generate** — type the new words you want that voice to say. Plain
-   text is enough; it is mapped to the reference speaker internally. You do not
-   need to write `S0:` or a speaker token.
-4. **Generate** — one click runs one synthesis. Requests are queued, and this
-   model accepts one generation at a time.
-5. **Generated audio** — play or download the result. The page returns a single
-   audio file.
-
-Example:
-
-- Reference transcript: `Hello, nice to meet you. My name is Chenyang.`
-- Text to generate: `Have a nice day and enjoy Southern California sunshine.`
-
-The output should say the new line in the reference voice, not repeat the
-introduction.
-
-This page is a single-speaker preview. Multi-speaker scenes, streaming, and the
-HTTP API are documented with SGLang Omni:
+# 🎭 Fish Audio Drama 3: direct a performance
+
+Write a line, pick a voice, and direct it. Drama 3 streams the take back as it is
+generated, and you can give the actor a note and get another take.
+
+**[fishaudio/drama-3](https://huggingface.co/fishaudio/drama-3)** (35B MoE, 3B active
+per token) runs on **one dedicated GPU** under
+**[SGLang Omni](https://github.com/sgl-project/sglang-omni)**, the serving stack of the
+model card's benchmarks (on an H200: about 90 ms to first audio and 7× faster than real
+time for a single stream). Requests from different visitors are batched by the engine
+instead of queuing behind each other.
+
+The Space hosts the Gradio UI on **CPU Basic**. Generation runs on the authenticated
+**HuggingFaceM4/drama-3-h200** Inference Endpoint, on one H200 in GCP us-south1.
+UI rebuilds no longer download Drama 3 or restart its GPU server. The page shows
+H200 and the median time to first audio; it does not report the remote GPU's memory.
+
+## What you can do
+
+- **Clone** one of the three preset voices, upload a clip, or record with the
+  microphone. Clips are trimmed to 20 s at a pause and transcribed with Whisper, so
+  you never have to type the transcript (you can still edit it).
+- **Design** a voice from a text description, with no audio at all: *"gravelly old
+  male voice, about seventy, weathered and slow"*. **Keep this voice** turns a take
+  you like into a clone reference, so later lines keep that voice.
+- **Surprise**: no reference and no description; the model improvises a voice.
+- **Direct** with the controls of the technical report (Appendix C), in plain language:
+  point tags `[laughs]`, scoped tags `<whisper>…</whisper>`, style prefixes
+  `[icy calm] Fine.`, and director's notes (`Direction: …` / `Line: …`).
+- **Give a note** on a take (*"more amused than before; drop the hurt"*). The previous
+  takes go back into the prompt as earlier turns, so the note is read against what you
+  just heard.
+
+## Sign-in and quotas
+
+Guests get 5 takes per UTC day without signing in. Signing in with Hugging Face gives
+free accounts 20 more; Hugging Face PRO accounts, members of Team and Enterprise orgs,
+and members of `DEMO_UNLIMITED_ORGS` (HuggingFaceM4 and fishaudio) have no daily cap.
+A take counts once its audio starts, so refused requests are free. Every caller also
+has at most 2 takes running at once.
+
+Guests are counted by address and by a session cookie, whichever is higher. The Hub's
+proxy passes a client-supplied `X-Forwarded-For` through, so a determined guest can
+get around the guest allowance; the engine's global limits still bound the load.
+
+Scripts can call the API with a Hugging Face token in an `X-HF-Token` header (or as a
+Bearer token); the account's quota applies the same way.
+
+As in the faster-qwen3-tts demo, usage is counted under HMACs, never raw ids, in
+SQLite. Attach a Storage Bucket at `/data` to keep counts across restarts, and set
+`DEMO_USAGE_HASH_SECRET` as a Space secret so the keys stay stable; without them the
+counts reset when the Space restarts.
+
+## How it works
+
+```
+browser ──Gradio──▶ server.py (CPU, :7860) ──authenticated SSE──▶ H200 endpoint
+                  · queue + per-visitor limits        · preprocessing → engine → vocoder
+                  · Whisper turbo on CPU              · drama3_space_ext.py (prompt layouts)
+                  · voices and takes in memory
+```
+
+- `server.py` sends requests to the H200 endpoint and relays raw PCM to the browser as it arrives. The endpoint token is a server-side Space secret, never sent to the browser. The stage is a custom Gradio HTML component with the original waveform, take cards, and Web Audio player. Gradio streams a cumulative PCM packet ledger so browser updates cannot drop chunks; audio is queued on one playback clock. A completed WAV appears for replay after the live stream drains. The original full UI remains available at `/legacy`.
+- `drama3_space_ext.py` subclasses the engine's Drama 3 preprocessor. SGLang Omni
+  does not let a stage factory be overridden from the command line, so `Dockerfile.engine`
+  appends one line to its `fishaudio_drama3/stages.py` (pinned commit) that rebinds the
+  class the factory builds. The stock `/v1/audio/speech` builds one prompt
+  layout: references, then one line. Drama 3 was also trained on casting a character
+  from a description (Figure 13 of the report) and on multi-turn directing with the
+  earlier takes in context (Figure 12). The server sends those as a JSON envelope;
+  requests without one take the stock path unchanged.
+- Sampling defaults are temperature 1.0, top-p 0.9, and no top-k cap (the server
+  sends top-k 4097: end-of-speech plus all 4,096 semantic codes). The SGLang Omni
+  default of top-k 30 makes Drama 3 fall silent mid-take: on the H200 endpoint, 40
+  takes per setting across the comparison scripts left 11 silent to the length cap
+  at top-k 30 (22% word error) and none without the cap (0.5% word error). At a
+  pause a few silence codes outrank the many ways speech can resume, and top-k
+  keeps only the former. The audio limit scales generously with the spoken script (at least 18 s,
+  with extra room for slow directions), up to about 56 s. A take that reaches the
+  limit is marked as possibly incomplete.
+- Drama 3 sometimes falls silent mid-line or fails to end the take. The server
+  trims leading silence, keeps at most 0.75 s of internal silence and 0.15 s at the
+  end, using the existing RMS threshold of 0.003. Four seconds of continuous
+  generated silence aborts the request. The first 2.5 s of audio is buffered (short
+  completed lines play immediately on completion), so an empty take or a stalled
+  opening can be discarded and retried once before playback. Failed openings do
+  not count against the daily allowance. This adds a little time before playback
+  in exchange for avoiding early broken takes.
+- Completed SSE events close the engine stream immediately. Without audio, a request
+  times out after 120 s before the first chunk or 15 s between chunks; each remote take also
+  has a 180 s wall-clock limit once the endpoint is ready. Partial stalled or timed-out takes
+  stay playable, with an incomplete warning in both studio UIs. These safeguards
+  bound bad generations; they cannot guarantee that the model speaks every word.
+- The UI does not poll the endpoint while idle. A generation wakes a sleeping endpoint and waits up to 10 minutes, with progress updates; this preserves its 60-minute scale-to-zero setting. Local GPU mode still warms cloned and unreferenced takes before enabling the page.
+
+## Gradio studio and Workflow
+
+The main page uses native Gradio controls, styled to match the original warm neutral
+palette, rounded cards, pill buttons, and colored direction tags. It keeps all script
+and character examples, Clone / Design / Surprise, voice upload and recording,
+automatic transcription, streaming playback and downloads, sampling, take history,
+retakes, and Keep this voice. Design opens with actual Sea captain values, ready to perform without typing. The stage keeps the original player appearance inside `gr.HTML`, with Gradio callbacks for generation, notes, and keeping voices. Cmd/Ctrl+Enter performs a take; the circle button switches
+between light and dark themes.
+
+The **Workflow** tab embeds a real `gr.Workflow` app at `/workflow/`. Its starter graph
+connects the script, voice mode (`Clone`, `Design`, or `Surprise`), preset voice id
+(`clone1`, `clone2`, or `clone3`), and voice description to the same generation backend,
+with an audio output. Both views use the existing quotas and admission limits.
+Workflow calls run on the server's event loop so its engine connections and queue are
+shared correctly. The studio exposes `/perform`, `/retake`, and `/keep_voice` through
+the Gradio API; the workflow exposes `/performance`.
+
+Gradio requires Workflow to be created at the top level, so it is mounted separately
+and embedded in a tab instead of constructed inside `gr.Blocks`. The parent app signs
+one shared OAuth session. Owners can edit the graph on the canvas after signing in;
+other visitors can run it. Workflow edits change the canvas/API; the custom Studio layout stays fixed.
+The local Drama 3 node works with the existing OAuth permissions. Adding remote
+inference nodes or using Save to Space requires enabling the corresponding
+`inference-api` or `write-repos` OAuth scopes separately.
+
+The editable graph is seeded once from `workflow.json` into a writable runtime path.
+Set `WORKFLOW_GRAPH=/data/drama3-workflow.json` to persist edits in the existing storage
+bucket, or leave it unset to use `/tmp/drama3-gradio/workflow.json`. Audio exports are
+kept in the Gradio temp directory, outside the application source. The endpoint keeps its pinned GPU image and SGLang Omni version; Gradio is pinned to **6.28.0**. `Dockerfile` is a small CPU UI image; `Dockerfile.engine` retains the original GPU recipe. Whisper turbo remains available for automatic transcription through faster-whisper int8 on CPU, cached in a Docker layer before the UI source.
+
+To preview the UI on a machine without a GPU (no model download or engine launch):
+
+```bash
+python -m venv .venv
+.venv/bin/pip install -r requirements-preview.txt
+DEMO_PREVIEW=1 .venv/bin/python server.py
+```
+
+Open `http://localhost:7860`. The preview loads the real preset voices and supports the
+UI controls; generation explicitly reports that the GPU engine is offline.
+Do not set `DEMO_PREVIEW` for the deployed Space. The original page is retained at
+`/legacy` for visual comparison and rollback.
+
+The adapter tests use a simulated engine stream against the real backend logic to
+check streaming, cancellation, quota failures, retakes, and Workflow execution:
+
+```bash
+.venv/bin/pip install pytest
+DEMO_PREVIEW=1 .venv/bin/python -m pytest tests -q
+```
+
+## Setup for fishaudio/Drama-3
+
+This app replaces the previous in-process ZeroGPU demo with the Studio and Workflow
+from [andito/Drama-3](https://huggingface.co/spaces/andito/Drama-3). Speech generation
+uses the existing authenticated **HuggingFaceM4/drama-3-h200** endpoint; its traffic
+and endpoint billing remain with HuggingFaceM4.
+
+After merging:
+
+1. In **Space Settings → Hardware**, switch from ZeroGPU to **CPU Basic**.
+2. In **Space Settings → Variables and secrets**, add **`DRAMA_ENDPOINT_TOKEN` as a
+   Secret**. Use the dedicated fine-grained token with permission to make calls to
+   Inference Endpoints in **HuggingFaceM4**. Repository access, endpoint management,
+   and Inference Providers permissions are not needed.
+3. Restart the Space after configuring the hardware and Secret, if it has not
+   already restarted automatically.
+
+The endpoint URL is already configured in `Dockerfile`; setting a
+`DRAMA_ENDPOINT_URL` variable is optional. The README's `hf_oauth: true` enables HF
+sign-in, and Hugging Face supplies its OAuth credentials automatically. No manual
+OAuth client or GPU model preload is required. The endpoint token stays on the
+server and is never sent to the browser.
+
+No persistent storage is required. With the default temporary storage, usage counts
+and saved Workflow edits reset when the Space restarts. Storage and a stable
+`DEMO_USAGE_HASH_SECRET` are optional settings described below.
+
+## Settings
+
+The deployed UI requires `DRAMA_ENDPOINT_TOKEN` as a **Space secret** with access to the endpoint. The URL is configured in the CPU Docker image and can be overridden with `DRAMA_ENDPOINT_URL`. Other settings are optional:
+
+| Variable | Default | What it does |
+|---|---|---|
+| `DRAMA_ENDPOINT_URL` | H200 endpoint URL | Remote speech server; omit outside the CPU Docker image for local GPU mode |
+| `DRAMA_ENDPOINT_TOKEN` | | Server-side authentication secret for the endpoint |
+| `ENDPOINT_WAKE_TIMEOUT_S` | `600` | Maximum wait for a sleeping endpoint |
+| `WHISPER_MODEL` | `turbo` | Automatic reference transcription model |
+| `WHISPER_CPU_THREADS` | `2` | CPU transcription threads |
+| `DEMO_MAX_ACTIVE` | `12` | Streams sent to the engine at once; the rest queue |
+| `DEMO_MAX_QUEUE` | `64` | Queue length before new requests are turned away |
+| `DEMO_GUEST_DAILY_REQUESTS` | `5` | Takes per UTC day without signing in |
+| `DEMO_DAILY_FREE_REQUESTS` | `20` | Takes per UTC day for signed-in free accounts |
+| `DEMO_REQUIRE_LOGIN` | `0` | Set to `1` to require sign-in for everything |
+| `DEMO_UNLIMITED_ORGS` | `fishaudio` | Orgs without a daily cap, besides HuggingFaceM4 |
+| `DEMO_USAGE_HASH_SECRET` | | Secret for the account keys in the usage database |
+| `DEMO_PER_CLIENT_ACTIVE` | `2` | Concurrent takes per account |
+| `DEMO_PER_CLIENT_REQUESTS` | `60` | Takes per account per `DEMO_RATE_WINDOW_S` (600 s), PRO included |
+| `ENGINE_MAX_RUNNING` | `16` | `--tts_engine.factory.max_running_requests` |
+| `ENGINE_EXTRA_ARGS` | | Extra `sgl-omni serve` flags, e.g. `--tts_engine.gpu_memory_fraction 0.8` |
+
+The CPU UI does not preload the speech model. The endpoint owns its weights and CUDA graphs. See [endpoint/README.md](endpoint/README.md) for the deployment scripts, pinned model/image, and native API. To run the model inside a GPU container instead, build with `docker build -f Dockerfile.engine .` and leave `DRAMA_ENDPOINT_URL` unset.
+
+## Run the full model yourself
 
 - [Fish Audio Drama 3 cookbook](https://sgl-project.github.io/sglang-omni/cookbook/fishaudio_drama3.html): install, voice cloning, and the HTTP API
 - [SGLang Omni](https://github.com/sgl-project/sglang-omni): the GitHub repo — please give it a star
-- [Technical report](https://github.com/fishaudio/fish-speech/tree/main/docs/assets/Fish_Audio_Drama_3_Technical_Report.pdf): architecture of the 35B MoE model
+- [Technical report](https://arxiv.org/abs/2610.18010): architecture of the 35B MoE model and its control syntax
 - [Fish Audio Playground](https://fish.audio): try Fish Audio voices in the browser
 - [Fish Speech](https://github.com/fishaudio/fish-speech): source repository for the model
 
-Sampling is fixed: `temperature=1.0`, `top_p=0.9`, `top_k=30`, and
-`max_new_tokens=500` per chunk. The token budget is also capped by the context
-window, so a chunk never asks for more frames than the prompt leaves free.
-Long text is split into chunks inside that single generation call.
-
-Each click books its own GPU slot with a dynamic `@spaces.GPU(duration=...)`
-estimate (≤120 s requested; ZeroGPU books 1.5× on this hardware). GPU time
-counts against the visitor's daily ZeroGPU quota — sign in for a larger
-allowance.
-
 ## Acknowledgments
 
 It took us over half a year to bring SGLang Omni back to life. Thanks to all the people who helped us along the way.
 
-Thanks so much for [AK](https://x.com/_akhaliq) help to set up the space.
+Thanks so much to [andito](https://huggingface.co/andito), [AK](https://x.com/_akhaliq), and [poli](https://huggingface.co/multimodalart) for helping set up the Space.
 
 **SGLang Omni team:** [Chenyang Zhao](https://x.com/GenAI_is_real), Jingwen Gu, Shenggui Li, Richard Wang, Yichi Zhang, Xuesong Ye, Jiaxin Deng, Ratish P, Jiaxuan Luo, Xinyu Lu, Wenyao Gao, Ao Sun, Jintao Quaige Li, Yuhao Chen, Yuhan Zhou, Xuxiang Tian, Ruoyi Zhang, Junnan Li, Yixiao Zeng, Chenchen Hong, Huapeng Zhou, Yang Gao, Chaoye Fei, Yiyang Zh, Kang, Yitian Gong, Yuan Luo, Hao Jin, Yifei Gao, Gaokai Zhang, Lifan Shen, Hui Shen, Kejing Xia, Jinxin Li, Dongming Shen, Jun Liu, Yue Yin, Zhihao Guo, Jisen Li, Ruoyu He, Xuanyi Li, Chenye Zhu, Liang Zhen, Xiaodong Ye, Yizhi Zhen, Kaizuo Li, Chonghan Liu, Wei Guan, Zhuofeng Li
 
 **FishAudio Team:** Shijia Liao, Shidong Li, Xuehao Yang, Jiahua Liu
 
-## Notes
+## License
 
-- Weights are preloaded at build time via `preload_from_hub`; first boot still
-  takes a while while the ~40 GB checkpoint is packed.
-- Model license: [Fish Audio Community License](https://huggingface.co/fishaudio/drama-3/blob/main/LICENSE)
-  — free non-commercial/research demo. Built with Fish Audio.
+Model license: [Fish Audio Community License](https://huggingface.co/fishaudio/drama-3/blob/main/LICENSE)
+— free non-commercial/research demo. Built with Fish Audio. Only clone voices you have
+permission to use.

accounts.py

@@ -0,0 +1,262 @@
+"""Who is asking, and how many takes they have left today.
+
+Guests can perform DEMO_GUEST_DAILY_REQUESTS takes per UTC day without signing in.
+Signing in with Hugging Face gives free accounts DEMO_DAILY_FREE_REQUESTS more, and
+PRO accounts, members of Team or Enterprise orgs, and members of the allow-listed
+orgs have no daily cap (entitlements.py).
+
+As in the faster-qwen3-tts demo, usage is stored under HMACs, never raw ids, in
+SQLite on an attached storage bucket (USAGE_BUCKET_MOUNT_PATH, /data) when one is
+mounted, else in /tmp. A signed-in account is keyed by its stable id. A guest is
+keyed by both its address and a session cookie, and the higher count applies; the
+Hub's proxy passes a client-supplied X-Forwarded-For through, so a determined guest
+can dodge this, which is why the engine's own concurrency limits stay in place.
+
+Scripts can identify with a Hugging Face token (X-HF-Token, or a Bearer token)
+instead of the browser session; the account's quota applies the same way.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import hmac
+import os
+import secrets
+import sqlite3
+import threading
+import time
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from types import SimpleNamespace
+
+from fastapi import HTTPException, Request
+from huggingface_hub import parse_huggingface_oauth
+
+from entitlements import Entitlement, _whoami_via_token, resolve_entitlement
+
+
+def env_flag(name: str, default: str = "0") -> bool:
+    return os.environ.get(name, default).strip().lower() not in {"0", "false", "no", "off", ""}
+
+
+REQUIRE_LOGIN = env_flag("DEMO_REQUIRE_LOGIN", "0")
+GUEST_DAILY_REQUESTS = int(os.environ.get("DEMO_GUEST_DAILY_REQUESTS", "5"))
+DAILY_FREE_REQUESTS = int(os.environ.get("DEMO_DAILY_FREE_REQUESTS", "20"))
+USAGE_BUCKET_MOUNT_PATH = Path(os.environ.get("USAGE_BUCKET_MOUNT_PATH", "/data"))
+USAGE_DB_FILENAME = os.environ.get("USAGE_DB_FILENAME", "drama3-usage.sqlite3")
+
+
+def default_usage_db_path() -> Path:
+    if USAGE_BUCKET_MOUNT_PATH.exists() and os.access(USAGE_BUCKET_MOUNT_PATH, os.W_OK):
+        return USAGE_BUCKET_MOUNT_PATH / USAGE_DB_FILENAME
+    return Path("/tmp/drama3-space") / "usage.sqlite3"
+
+
+USAGE_DB_PATH = Path(os.environ.get("USAGE_DB_PATH", str(default_usage_db_path())))
+# Set DEMO_USAGE_HASH_SECRET as a Space secret when usage persists in a bucket;
+# without it, keys stay private but change on every restart.
+_HASH_SECRET = (os.environ.get("DEMO_USAGE_HASH_SECRET") or secrets.token_hex(32)).encode()
+_db_lock = threading.Lock()
+_db_ready = False
+
+GUEST = Entitlement(tier="guest", is_pro=False, is_team_member=False)
+LOCAL = Entitlement(tier="team", is_pro=False, is_team_member=True)
+
+
+@dataclass(frozen=True)
+class Identity:
+    """keys[0] identifies the caller for concurrency limits; usage counts all keys."""
+
+    keys: tuple[str, ...]
+    name: str
+    entitlement: Entitlement
+    signed_in: bool
+
+    @property
+    def key(self) -> str:
+        return self.keys[0]
+
+    @property
+    def daily_limit(self) -> int | None:
+        if self.entitlement.unlimited:
+            return None
+        return DAILY_FREE_REQUESTS if self.signed_in else GUEST_DAILY_REQUESTS
+
+
+def hashed(kind: str, value: str) -> str:
+    digest = hmac.new(_HASH_SECRET, f"{kind}:{value}".encode("utf-8"), hashlib.sha256).hexdigest()
+    return f"{kind}_{digest}"
+
+
+def request_token(request: Request) -> str | None:
+    token = request.headers.get("x-hf-token", "").strip()
+    if not token:
+        scheme, _, value = request.headers.get("authorization", "").partition(" ")
+        token = value.strip() if scheme.lower() == "bearer" else ""
+    return token if token.startswith("hf_") else None
+
+
+def signed_in_identity(request: Request) -> Identity | None:
+    """The account behind the browser session or a Hugging Face token, if any."""
+    oauth_info = parse_huggingface_oauth(request)
+    user = getattr(oauth_info, "user_info", None) if oauth_info is not None else None
+    if user is None:
+        token = request_token(request)
+        profile = _whoami_via_token(token) if token else {}
+        if not profile.get("id"):
+            return None
+        user = SimpleNamespace(
+            sub=profile["id"],
+            preferred_username=profile.get("name"),
+            is_pro=profile.get("isPro", False),
+            orgs=profile.get("orgs", []),
+        )
+        oauth_info = SimpleNamespace(user_info=user, access_token=token)
+    account_id = str(getattr(user, "sub", None) or getattr(user, "preferred_username", "") or "")
+    if not account_id:
+        return None
+    name = str(getattr(user, "preferred_username", None) or getattr(user, "name", None) or account_id)
+    return Identity(
+        keys=(hashed("hfuser", account_id),),
+        name=name,
+        entitlement=resolve_entitlement(oauth_info),
+        signed_in=True,
+    )
+
+
+def guest_identity(request: Request) -> Identity:
+    forwarded = request.headers.get("x-forwarded-for", "")
+    address = forwarded.split(",", 1)[0].strip() if forwarded else (request.client.host if request.client else "")
+    keys = [hashed("guestip", address or "unknown")]
+    try:
+        session = request.session
+    except AssertionError:  # no session middleware, e.g. when run outside a Space
+        session = None
+    if session is not None:
+        session.setdefault("guest", secrets.token_urlsafe(18))
+        keys.append(hashed("guest", session["guest"]))
+    return Identity(keys=tuple(keys), name="guest", entitlement=GUEST, signed_in=False)
+
+
+def identify(request: Request) -> Identity:
+    identity = signed_in_identity(request)
+    if identity is not None:
+        return identity
+    if REQUIRE_LOGIN:
+        raise HTTPException(401, "Sign in with Hugging Face to use this demo.")
+    if not os.environ.get("SPACE_ID"):
+        # Running locally: no quotas.
+        return Identity(keys=("local",), name="local", entitlement=LOCAL, signed_in=False)
+    return guest_identity(request)
+
+
+# ─── Daily usage ──────────────────────────────────────────────────────────────
+
+
+def today() -> str:
+    return datetime.now(timezone.utc).date().isoformat()
+
+
+def _connect() -> sqlite3.Connection:
+    global _db_ready
+    if not _db_ready:
+        USAGE_DB_PATH.parent.mkdir(parents=True, exist_ok=True)
+    connection = sqlite3.connect(USAGE_DB_PATH, timeout=30)
+    if not _db_ready:
+        connection.execute(
+            """CREATE TABLE IF NOT EXISTS usage_daily (
+                user_key TEXT NOT NULL, day TEXT NOT NULL,
+                is_pro INTEGER NOT NULL DEFAULT 0, count INTEGER NOT NULL DEFAULT 0,
+                updated_at INTEGER NOT NULL, PRIMARY KEY (user_key, day))"""
+        )
+        connection.execute(
+            """CREATE TABLE IF NOT EXISTS usage_users (
+                user_key TEXT PRIMARY KEY, username TEXT NOT NULL,
+                is_pro INTEGER NOT NULL DEFAULT 0,
+                first_seen_at INTEGER NOT NULL, last_seen_at INTEGER NOT NULL)"""
+        )
+        _db_ready = True
+    return connection
+
+
+def _count(connection: sqlite3.Connection, identity: Identity, day: str) -> int:
+    counts = [0]
+    for key in identity.keys:
+        row = connection.execute(
+            "SELECT count FROM usage_daily WHERE user_key = ? AND day = ?", (key, day)
+        ).fetchone()
+        counts.append(int(row[0]) if row else 0)
+    return max(counts)
+
+
+def _record_user(connection: sqlite3.Connection, identity: Identity, now: int) -> None:
+    if not identity.signed_in:
+        return
+    connection.execute(
+        """INSERT INTO usage_users (user_key, username, is_pro, first_seen_at, last_seen_at)
+           VALUES (?, ?, ?, ?, ?)
+           ON CONFLICT(user_key) DO UPDATE SET username = excluded.username,
+               is_pro = excluded.is_pro, last_seen_at = excluded.last_seen_at""",
+        (identity.key, identity.name, int(identity.entitlement.is_pro), now, now),
+    )
+
+
+def usage_payload(identity: Identity, count: int) -> dict:
+    limit = identity.daily_limit
+    return {
+        "used_today": count,
+        "limit": limit,
+        "remaining": None if limit is None else max(0, limit - count),
+        "tier": identity.entitlement.tier,
+        "unlimited": limit is None,
+        "signed_in": identity.signed_in,
+        "free_after_sign_in": DAILY_FREE_REQUESTS,
+    }
+
+
+def usage(identity: Identity) -> dict:
+    with _db_lock, _connect() as connection:
+        _record_user(connection, identity, int(time.time()))
+        return usage_payload(identity, _count(connection, identity, today()))
+
+
+def check_quota(identity: Identity) -> None:
+    limit = identity.daily_limit
+    if limit is None:
+        return
+    with _db_lock, _connect() as connection:
+        if _count(connection, identity, today()) < limit:
+            return
+    if identity.signed_in:
+        raise HTTPException(
+            429,
+            f"That's your {limit} free takes for today. Hugging Face PRO accounts have "
+            "unlimited takes; the free ones come back at midnight UTC.",
+        )
+    raise HTTPException(
+        429,
+        f"That's the {limit} takes guests get each day. Sign in with Hugging Face for "
+        f"{DAILY_FREE_REQUESTS} more.",
+    )
+
+
+def consume_take(identity: Identity) -> dict:
+    """Count one take for today under every key, and return the updated usage."""
+    day, now = today(), int(time.time())
+    with _db_lock, _connect() as connection:
+        _record_user(connection, identity, now)
+        count = _count(connection, identity, day) + 1
+        for key in identity.keys:
+            connection.execute(
+                """INSERT INTO usage_daily (user_key, day, is_pro, count, updated_at)
+                   VALUES (?, ?, ?, ?, ?)
+                   ON CONFLICT(user_key, day) DO UPDATE SET is_pro = excluded.is_pro,
+                       count = excluded.count, updated_at = excluded.updated_at""",
+                (key, day, int(identity.entitlement.is_pro), count, now),
+            )
+    return usage_payload(identity, count)
+
+
+def public_user(identity: Identity) -> dict:
+    return {"username": identity.name, "tier": identity.entitlement.tier, "signed_in": identity.signed_in}

app.py

@@ -1,776 +0,0 @@
-"""Single-speaker Drama 3 voice cloning demo for Hugging Face ZeroGPU.
-
-Inputs: reference audio, its exact transcript, and the text to synthesize.
-The UI is the fishaudio team's hand-built Blocks page, mounted inside the app
-view of a `gr.Workflow` — the built-in toggle switches between this page and
-the Workflow canvas (workflow.json), which runs the same single-speaker
-pipeline through the `drama3_clone` node. Keep fish_speech/, requirements.txt,
-and the README SDK configuration.
-
-Inference is the official PyTorch path of fishaudio/drama-3, vendored under
-fish_speech/. It runs single-process on a ZeroGPU `large` slice (RTX Pro 6000,
-48 GB): the 32.2B FP8 MoE experts stay in FP8 — dequantised to bf16 they would
-need ~67 GB. Modules are built on the meta device, materialised on CPU with FP8
-experts, then moved to "cuda", because the FP8 path needs an SM89+ GPU at load
-time and ZeroGPU has no GPU at import.
-
-Requires the Space hardware to be set to ZeroGPU (Settings -> Hardware).
-"""
-
-import os
-
-os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
-
-# Import spaces before any module initializes CUDA. Model loading stays at
-# module scope so ZeroGPU can pack the weights during startup.
-try:
-    import spaces
-except ImportError:
-    class _SpacesShim:
-        @staticmethod
-        def GPU(duration=60):
-            return lambda fn: fn
-
-    spaces = _SpacesShim()
-
-import difflib
-import json
-import random
-import re
-import time
-import urllib.parse
-from pathlib import Path
-
-import gradio as gr
-import librosa
-import numpy as np
-import soundfile as sf
-import torch
-from gradio_client import utils as client_utils
-from tokenizers import Tokenizer
-
-from fish_speech.codec import decode as codec_decode
-from fish_speech.codec import encode as codec_encode
-from fish_speech.codec import load_codec
-from fish_speech.models.text2semantic.errors import InvalidInputError
-from fish_speech.models.text2semantic.fish_omni.backbone import (
-    CHECKPOINT_BACKBONE_PREFIX,
-    CHECKPOINT_LM_HEAD_KEY,
-    FishOmniBackbone,
-    load_backbone_weights,
-)
-from fish_speech.models.text2semantic.fish_omni.chunking import (
-    Drama3History,
-    split_drama3_text,
-)
-from fish_speech.models.text2semantic.fish_omni.config import FishOmniConfig
-from fish_speech.models.text2semantic.fish_omni.generate import (
-    Drama3DecodeGraphs,
-    Drama3PromptTooLongError,
-    build_drama3_prompt,
-    generate_drama3,
-)
-from fish_speech.models.text2semantic.fish_omni.model import (
-    AUDIO_DECODER_PREFIXES,
-    FishOmniModel,
-    iter_checkpoint_tensors,
-)
-from fish_speech.models.text2semantic.fish_omni.tower import (
-    FishOmniAudioDecoder,
-    load_audio_decoder_weights,
-)
-
-MODEL_ID = "fishaudio/drama-3"
-DTYPE = torch.bfloat16
-SAMPLE_RATE = 44100
-MAX_REF_SECONDS = 30
-MAX_SCRIPT_CHARS = 3000
-MAX_SEED = 2**31 - 1
-# Sampling values used by the SGLang-Omni Drama 3 Cookbook's serving examples.
-# These are model-generation parameters, independent of the UI/runtime backend.
-TEMPERATURE = 1.0
-TOP_P = 0.9
-TOP_K = 30
-MAX_NEW_TOKENS = 500
-COOKBOOK_URL = "https://sgl-project.github.io/sglang-omni/cookbook/fishaudio_drama3.html"
-SGLANG_OMNI_URL = "https://github.com/sgl-project/sglang-omni"
-TECHNICAL_REPORT_URL = (
-    "https://github.com/fishaudio/fish-speech/tree/main/docs/assets/"
-    "Fish_Audio_Drama_3_Technical_Report.pdf"
-)
-FISH_AUDIO_URL = "https://fish.audio"
-FISH_SPEECH_URL = "https://github.com/fishaudio/fish-speech"
-MODEL_URL = "https://huggingface.co/fishaudio/drama-3"
-AK_URL = "https://x.com/_akhaliq"
-CHENYANG_URL = "https://x.com/GenAI_is_real"
-SHORT_SPEAKER = re.compile(r"^\s*(?:\[S(\d+)\]|S(\d+)\s*:)\s*", re.M)
-SPEAKER_TAG = re.compile(r"<\|speaker:(\d+)\|>")
-
-
-def speaker_tag(index=0):
-    return "<" + f"|speaker:{index}|>"
-
-
-# Preserve the original FP8 model-loading path. On CPU the page opens without
-# attempting to load the checkpoint; generation requires ZeroGPU or CUDA.
-print(f"[load] {MODEL_ID} ...", flush=True)
-if torch.cuda.is_available():
-    from huggingface_hub import snapshot_download
-
-    checkpoint_dir = snapshot_download(MODEL_ID)
-    with open(os.path.join(checkpoint_dir, "config.json"), encoding="utf-8") as f:
-        config = FishOmniConfig.from_dict(json.load(f))
-    with torch.device("meta"):
-        backbone = FishOmniBackbone(
-            config.text, fp8_experts=config.fp8_experts, dtype=DTYPE
-        )
-        audio_decoder = FishOmniAudioDecoder(config, dtype=DTYPE)
-    backbone.materialize(["cpu"])
-    load_backbone_weights(
-        backbone,
-        iter_checkpoint_tensors(
-            Path(checkpoint_dir), skip_prefixes=AUDIO_DECODER_PREFIXES
-        ),
-    )
-    audio_decoder.to_empty(device="cpu")
-    load_audio_decoder_weights(
-        audio_decoder,
-        iter_checkpoint_tensors(
-            Path(checkpoint_dir),
-            skip_prefixes=(CHECKPOINT_BACKBONE_PREFIX, CHECKPOINT_LM_HEAD_KEY),
-        ),
-    )
-    tokenizer = Tokenizer.from_file(os.path.join(checkpoint_dir, "tokenizer.json"))
-    model = FishOmniModel(config, backbone, audio_decoder, tokenizer).eval()
-    model.to("cuda")
-    codec = load_codec(os.path.join(checkpoint_dir, "codec.pth"), DTYPE).to("cuda")
-else:
-    model = None
-    codec = None
-print("[load] ready", flush=True)
-
-
-# ---------------------------------------------------------------------------
-# Canvas value adapters (used by the workflow node only)
-# ---------------------------------------------------------------------------
-def to_audio_path(value):
-    """Normalize whatever the canvas audio port sends to a local path (or None).
-
-    Canvas file values carry only a url (FileValue has no path): same-server
-    /gradio_api/file= urls resolve to local paths, http(s) urls are fetched.
-    (sample_rate, array) tuples are written to a temp wav.
-    """
-    if value is None:
-        return None
-    if isinstance(value, tuple) and len(value) == 2:
-        from gradio.utils import get_upload_folder
-
-        sr, data = value
-        directory = get_upload_folder()
-        os.makedirs(directory, exist_ok=True)
-        path = os.path.join(directory, f"in_{os.urandom(8).hex()}.wav")
-        sf.write(path, np.asarray(data), int(sr))
-        return path
-    if isinstance(value, str):
-        src = value
-    elif isinstance(value, dict):
-        src = value.get("path") or value.get("url")
-    else:
-        return None
-    if not isinstance(src, str) or not src:
-        return None
-    if src.startswith("/gradio_api/file="):
-        src = urllib.parse.unquote(src.removeprefix("/gradio_api/file="))
-    if src.startswith(("http://", "https://")):
-        import httpx
-
-        from gradio.utils import get_upload_folder
-
-        resp = httpx.get(src, timeout=60, follow_redirects=True)
-        resp.raise_for_status()
-        directory = get_upload_folder()
-        os.makedirs(directory, exist_ok=True)
-        path = os.path.join(directory, f"in_{os.urandom(8).hex()}.wav")
-        with open(path, "wb") as f:
-            f.write(resp.content)
-        return path
-    return src if os.path.exists(src) else None
-
-
-def to_audio_media(sample_rate, audio):
-    """Save an int16 waveform and return the media dict the canvas renders.
-
-    Saved into gradio's upload folder (as gradio's own workflow code does),
-    which is served without an allowlist entry.
-    """
-    from gradio.utils import get_upload_folder
-
-    directory = get_upload_folder()
-    os.makedirs(directory, exist_ok=True)
-    path = os.path.join(directory, f"drama3_{os.urandom(8).hex()}.wav")
-    sf.write(path, audio, sample_rate)
-    return {
-        "path": path,
-        "url": f"/gradio_api/file={client_utils.encode_file_path(path)}",
-        "is_file": True,
-    }
-
-
-# ---------------------------------------------------------------------------
-# Demo logic (unchanged from the fishaudio team's single-speaker version)
-# ---------------------------------------------------------------------------
-def normalize_script(text):
-    """Accept plain text and bind it to the single reference speaker."""
-    text = (text or "").replace("\r\n", "\n").strip()
-    if not text:
-        raise gr.Error("Enter the text you want to generate.")
-    text = SHORT_SPEAKER.sub(
-        lambda m: speaker_tag(m.group(1) or m.group(2)) + " ", text
-    )
-    if any(int(index) != 0 for index in SPEAKER_TAG.findall(text)):
-        raise gr.Error("This demo supports one voice. Enter a single speaker's text.")
-    if not text.startswith(speaker_tag()):
-        text = speaker_tag() + " " + text
-    if len(text) > MAX_SCRIPT_CHARS:
-        raise gr.Error(f"Keep the text under {MAX_SCRIPT_CHARS} characters.")
-    return text
-
-
-def load_reference(path):
-    """Validate the entire clip without silently dropping transcript content."""
-    try:
-        wav, _ = librosa.load(path, sr=SAMPLE_RATE, mono=True)
-    except Exception as err:
-        raise gr.Error("Could not read the reference audio. Try uploading a WAV file.") from err
-    if len(wav) < SAMPLE_RATE // 2:
-        raise gr.Error("Record or upload at least half a second of reference speech.")
-    if len(wav) > MAX_REF_SECONDS * SAMPLE_RATE:
-        raise gr.Error("Use a reference clip of 30 seconds or less, with its matching transcript.")
-    if not np.isfinite(wav).all():
-        raise gr.Error("The reference audio contains invalid samples. Use another recording.")
-    peak = np.abs(wav).max()
-    if peak == 0:
-        raise gr.Error("The reference audio is silent. Record a spoken sentence.")
-    wav = wav * min(1.0, 0.95 / peak)
-    return torch.from_numpy(wav.astype(np.float32))
-
-
-def build_chunk_prompt(text, ref_texts, ref_codes, history):
-    error = None
-    for turns in history.candidates():
-        try:
-            return build_drama3_prompt(
-                text, tokenizer, config, ref_texts, ref_codes, history=turns
-            )
-        except Drama3PromptTooLongError as too_long:
-            error = too_long
-    if error is not None:
-        raise error
-    raise gr.Error("Could not build the speech prompt. Try shorter text and audio.")
-
-
-def estimate_duration(script, reference_waveform, reference_text):
-    # Keep the original duration estimate and 120-second request cap.
-    chunks = max(1, -(-len(script.encode()) // 300))
-    return int(min(120, 8 + chunks * 19 + 4))
-
-
-@spaces.GPU(duration=estimate_duration)
-def _gpu_synthesize(script, reference_waveform, reference_text):
-    """One synthesis call; the reference is conditioning, never a prior take."""
-    started = time.perf_counter()
-    with torch.inference_mode():
-        seed = int(np.random.randint(0, MAX_SEED))
-        torch.manual_seed(seed)
-        np.random.seed(seed)
-        ref_texts = [speaker_tag() + " " + reference_text]
-        ref_codes = [codec_encode(codec, reference_waveform).cpu()]
-        chunks = split_drama3_text(script) or [script]
-        capacity = Drama3DecodeGraphs.graph_capacity(model)
-        history = Drama3History()
-        pieces = []
-        finished_all = True
-        try:
-            for chunk in chunks:
-                prompt = build_chunk_prompt(chunk, ref_texts, ref_codes, history)
-                budget = min(MAX_NEW_TOKENS, capacity - len(prompt))
-                if budget < 1:
-                    raise gr.Error("The reference and text are too long. Try a shorter clip.")
-                frames, finish = generate_drama3(
-                    model,
-                    prompt,
-                    max_new_tokens=budget,
-                    temperature=TEMPERATURE,
-                    top_k=TOP_K,
-                    top_p=TOP_P,
-                )
-                finished_all &= finish == "eos"
-                if len(frames):
-                    codes = frames[:, 1:].T.contiguous()
-                    pieces.append(codes)
-                    history.push(chunk, codes)
-        except (InvalidInputError, Drama3PromptTooLongError) as err:
-            raise gr.Error(str(err)) from err
-        if not pieces:
-            raise gr.Error("No speech was generated. Click Generate to try again.")
-        codes = torch.cat(pieces, dim=1)
-        audio = codec_decode(codec, codes).float().cpu().numpy()
-        audio = np.clip(audio, -1.0, 1.0)
-        if not finished_all:
-            gr.Warning("The generated audio may be cut short. Try shorter text.")
-        elapsed = time.perf_counter() - started
-        seconds = len(audio) / SAMPLE_RATE
-        info = (
-            f"seed {seed} | {seconds:.1f}s of audio in {elapsed:.1f}s | "
-            f"chunks: {len(chunks)}"
-        )
-        return (SAMPLE_RATE, (audio * 32767).astype(np.int16)), seed, info
-
-
-def synthesize(reference_audio, reference_text, text_to_generate, consent_verified):
-    # The consent gate is enforced here, not just in the UI.
-    if not consent_verified:
-        raise gr.Error(
-            "Complete the voice consent gate first: generate a sentence, record "
-            "yourself reading it, and verify."
-        )
-    # Validate on CPU before requesting scarce ZeroGPU time.
-    if not reference_audio:
-        raise gr.Error("Record or upload reference audio first.")
-    reference_text = (reference_text or "").strip()
-    if not reference_text:
-        raise gr.Error("Enter exactly what is spoken in the reference audio.")
-    script = normalize_script(text_to_generate)
-    if model is None:
-        raise gr.Error("Generation requires GPU hardware. Set this Space to ZeroGPU and restart it.")
-    waveform = load_reference(reference_audio)
-    (sr_audio, _, _info) = _gpu_synthesize(script, waveform, reference_text)
-    return sr_audio
-
-
-def drama3_clone(reference_audio, reference_text, text_to_generate, consent_confirmed=False):
-    """Workflow node: the same single-speaker pipeline, callable from the canvas.
-
-    Returns (audio media dict, seed, run-info text).
-    """
-    if not consent_confirmed:
-        raise gr.Error(
-            "Tick the consent box: only clone a voice whose speaker explicitly agreed."
-        )
-    if not reference_audio:
-        raise gr.Error("Record or upload reference audio first.")
-    reference_text = (reference_text or "").strip()
-    if not reference_text:
-        raise gr.Error("Enter exactly what is spoken in the reference audio.")
-    script = normalize_script(text_to_generate)
-    if model is None:
-        raise gr.Error("Generation requires GPU hardware. Set this Space to ZeroGPU and restart it.")
-    waveform = load_reference(to_audio_path(reference_audio) or reference_audio)
-    (sr, audio), seed, info = _gpu_synthesize(script, waveform, reference_text)
-    return to_audio_media(sr, audio), seed, info
-
-
-# ---------------------------------------------------------------------------
-# Voice consent gate (https://huggingface.co/blog/voice-consent-gate): a voice
-# is cloned only after the speaker reads a unique consent sentence aloud and an
-# ASR check recognizes it. The same recording then serves as the reference
-# audio — spoken consent and cloning input in one.
-# ---------------------------------------------------------------------------
-CONSENT_MODEL_NAME = "Drama 3"
-CONSENT_ASR_ID = "openai/whisper-base.en"
-CONSENT_THRESHOLD = 0.85
-CONSENT_BLOG_URL = "https://huggingface.co/blog/voice-consent-gate"
-LLAMA_SPACE_ID = os.getenv("LLAMA_SPACE_ID", "huggingface-projects/llama-3.2-3B-Instruct")
-LLAMA_API_NAME = "/chat"
-CONSENT_TOPICS = [
-    "the weather", "food", "music", "a morning walk", "coffee", "reading",
-    "gardening", "the weekend", "a favorite movie", "cooking dinner",
-    "public transport", "a local park",
-]
-# Used only if the sentence-generation Space is unreachable; the LLM path is
-# what makes each session's sentence novel.
-CONSENT_FALLBACK_SENTENCES = [
-    "I give my consent to use my voice for generating audio with the model Drama 3. The weather is bright and calm this morning.",
-    "I give my consent to use my voice for generating audio with the model Drama 3. The coffee shop outside has a pleasant aroma this morning.",
-    "I agree to the use of my recorded voice for audio generation with the model Drama 3. After a gentle walk, I am feeling relaxed and ready to speak.",
-    "I give my consent to use my voice for generating synthetic audio with Drama 3 today. My daily routine keeps me out and about most days.",
-]
-
-
-def get_consent_generation_prompt(topic: str) -> str:
-    return f"""Generate exactly two short, natural-sounding English sentences (10-15 words each) that a person could say aloud, using everyday language.
-
-Sentence 1 (Consent sentence):
-* Clearly states informed consent to use their voice for generating synthetic audio with an AI model called {CONSENT_MODEL_NAME}.
-* Must explicitly include a consent phrase such as "I give my consent," "I agree," or "I allow."
-* Must clearly mention the model name {CONSENT_MODEL_NAME} in the sentence.
-* Should sound fluent, polite, and natural to read aloud.
-
-Sentence 2 (Phonetic variety sentence):
-* Should not repeat the consent content.
-* Adds phonetic variety with a neutral descriptive clause, for example about {topic}.
-* Should be fluent, natural, and comfortable to read aloud, with diverse vowels and consonants.
-
-FORMAT:
-* Output EXACTLY two sentences.
-* No numbering, no quotes, no bullet points, and no introductory text.
-
-Example format (don't copy text, just the format):
-I give my consent to use my voice for generating audio with the model {CONSENT_MODEL_NAME}. The weather is clear and calm this afternoon, and I'm speaking at an even pace."""
-
-
-def generate_consent_sentence():
-    """A novel consent sentence pair per session (small LLM, local fallback)."""
-    topic = random.choice(CONSENT_TOPICS)
-    try:
-        from gradio_client import Client
-
-        client = Client(LLAMA_SPACE_ID)
-        result = client.predict(
-            message=get_consent_generation_prompt(topic),
-            max_new_tokens=128,
-            temperature=0.6,
-            top_p=0.9,
-            top_k=50,
-            repetition_penalty=1.2,
-            api_name=LLAMA_API_NAME,
-        )
-        text = " ".join(str(result).split())
-        if CONSENT_MODEL_NAME in text and len(text.split()) >= 8:
-            return text
-    except Exception as err:
-        print(f"[consent] LLM sentence generation failed, using fallback: {err!r}", flush=True)
-    return random.choice(CONSENT_FALLBACK_SENTENCES)
-
-
-_consent_asr = None
-
-
-def _load_consent_asr():
-    """Whisper on CPU; loaded on first verify, not at boot."""
-    global _consent_asr
-    if _consent_asr is None:
-        from transformers import pipeline as hf_pipeline
-
-        _consent_asr = hf_pipeline(
-            "automatic-speech-recognition", model=CONSENT_ASR_ID, device="cpu"
-        )
-    return _consent_asr
-
-
-def _word_tokens(text: str) -> list:
-    return re.sub(r"[^a-z0-9 ]", "", (text or "").lower()).split()
-
-
-def verify_consent(recording_path, sentence):
-    """Transcribe the consent recording; True when it matches the sentence."""
-    if not recording_path:
-        raise gr.Error("Record yourself reading the sentence first (microphone only).")
-    sentence = (sentence or "").strip()
-    if not sentence:
-        raise gr.Error("Generate a consent sentence first.")
-    wav, _ = librosa.load(recording_path, sr=16000, mono=True)
-    if len(wav) < 16000 // 2:
-        raise gr.Error("The recording is too short — read the whole sentence aloud.")
-    heard = _load_consent_asr()({"raw": wav, "sampling_rate": 16000})["text"].strip()
-    ratio = difflib.SequenceMatcher(
-        None, _word_tokens(heard), _word_tokens(sentence)
-    ).ratio()
-    return ratio >= CONSENT_THRESHOLD, heard, ratio
-
-
-def on_generate_sentence():
-    return (
-        generate_consent_sentence(),
-        False,
-        "Now read the sentence aloud into your microphone, then verify.",
-    )
-
-
-def on_verify(recording, sentence):
-    verified, heard, ratio = verify_consent(recording, sentence)
-    if not verified:
-        return {
-            consent_ok: False,
-            consent_status: (
-                f"❌ We heard: ��{heard}” ({ratio:.0%} match — needs "
-                f"{CONSENT_THRESHOLD:.0%}). Read the sentence exactly as shown "
-                "and verify again."
-            ),
-        }
-    return {
-        reference_audio: recording,
-        reference_text: sentence,
-        consent_ok: True,
-        consent_status: (
-            f"✅ Consent verified ({ratio:.0%} match). The recording is now your "
-            "reference audio, and the sentence its transcript."
-        ),
-    }
-
-
-# Fish Audio Drama 3 Space draft: cream page, white cards,
-# and a black pill button.
-FISH_CSS = """
-:root, .gradio-container {
-    --fa-bg: #f6f5f3; --fa-ink: #141414; --fa-muted: #6b6b66;
-    --fa-card: #ffffff; --fa-border: #e2e0dc; --fa-card-border: #e8e6e2;
-}
-.dark, .dark .gradio-container {
-    --fa-bg: #161615; --fa-ink: #f3f2ef; --fa-muted: #a8a7a1;
-    --fa-card: #201f1d; --fa-border: #35332f; --fa-card-border: #35332f;
-}
-:root, .gradio-container, .dark, .dark .gradio-container {
-    --body-text-color: var(--fa-ink);
-    --body-text-color-subdued: var(--fa-muted);
-    --block-label-text-color: var(--fa-ink);
-    --block-title-text-color: var(--fa-ink);
-    --button-primary-background-fill: var(--fa-ink);
-    --button-primary-text-color: var(--fa-bg);
-    --background-fill-primary: var(--fa-bg);
-    --background-fill-secondary: var(--fa-card);
-    --color-accent: var(--fa-ink);
-}
-body, .gradio-container, .app, .main {
-    background: var(--fa-bg) !important;
-    color: var(--fa-ink) !important;
-}
-.gradio-container .prose strong, .gradio-container .prose b,
-.gradio-container .prose h1, .gradio-container .prose h2, .gradio-container .prose h3,
-.gradio-container .prose li {
-    color: var(--fa-ink) !important;
-}
-.gradio-container {
-    max-width: 980px !important;
-    margin-left: auto !important;
-    margin-right: auto !important;
-    font-family: ui-sans-serif, -apple-system, "Segoe UI", Inter, Roboto, sans-serif !important;
-}
-#fa-hero { margin: 36px 0 8px; }
-.fa-pill {
-    display: inline-flex; align-items: center;
-    border: 1px solid var(--fa-border); border-radius: 999px;
-    background: var(--fa-card); color: var(--fa-ink);
-    font-size: 13px; font-weight: 500; padding: 5px 14px;
-}
-#fa-hero h1 {
-    font-size: clamp(30px, 5vw, 46px); font-weight: 800;
-    letter-spacing: -0.025em; line-height: 1.08;
-    max-width: 640px; margin: 18px 0 10px; color: var(--fa-ink);
-}
-#fa-hero p, #fa-intro p {
-    font-size: 15px; line-height: 1.55; color: var(--fa-muted);
-    max-width: 640px;
-}
-#fa-links { display: flex; gap: 8px; flex-wrap: wrap; margin: 18px 0 22px; }
-#fa-links a {
-    border: 1px solid var(--fa-border); border-radius: 999px;
-    background: var(--fa-card); padding: 6px 14px;
-    font-size: 13px; font-weight: 500;
-    color: var(--fa-ink) !important; text-decoration: none !important;
-}
-#fa-links a:hover { border-color: var(--fa-ink); }
-#fa-form .block, #fa-form .form {
-    border-radius: 16px !important; border-color: var(--fa-card-border) !important;
-    background: var(--fa-card) !important;
-    box-shadow: 0 1px 2px rgba(20,20,20,0.04) !important;
-}
-#fa-form .prose, #fa-guide .prose, #fa-ack .prose {
-    background: transparent !important; border: none !important;
-    box-shadow: none !important;
-}
-#fa-form .prose p { color: var(--fa-muted); font-size: 13.5px; line-height: 1.5; }
-#fa-form label span, #fa-form .label-wrap span {
-    font-weight: 600 !important; color: var(--fa-ink) !important;
-}
-#fa-form button.primary {
-    border-radius: 999px !important;
-    background: var(--fa-ink) !important; color: var(--fa-bg) !important;
-    font-weight: 600 !important; border: none !important;
-}
-#fa-guide h3, #fa-ack h3 { letter-spacing: -0.02em; }
-footer { display: none !important; }
-"""
-
-# Workflow additions only: hide the auto-generated form (the team's page below
-# replaces it) and dock the App/Workflow toggle bottom-center.
-WORKFLOW_CSS = """
-.workflow-app-body { display: none !important; }
-.view-toggle {
-    left: 50% !important; right: auto !important;
-    transform: translateX(-50%) !important;
-    bottom: 18px !important;
-}
-"""
-
-demo = gr.Workflow(
-    graph="workflow.json",
-    bind={"drama3_clone": drama3_clone},
-)
-
-# Mount the team's page inside the workflow's app view, in place of the
-# auto-generated form (hidden above). The canvas only covers #workflow-app's
-# overlay, so the toggle still switches to the workflow graph. Blocks appended
-# here are not tracked by WorkflowEndpointManager, so canvas saves keep them.
-from gradio.workflow_api import _active_blocks
-
-with _active_blocks(demo, demo._api_endpoints.app_root):
-    gr.HTML(
-        """
-        <div id="fa-hero">
-          <span class="fa-pill">Try · Drama 3 →</span>
-          <h1>The most expressive, emotionally controllable real-time voice model</h1>
-          <p>Five seconds is easy. Five minutes is the test. Drama 3 clones a voice
-          from a few seconds of reference audio and holds its character from the
-          first word to the last.</p>
-        </div>
-        """
-    )
-    gr.Markdown(
-        "Clone one voice, then speak a new line in that voice. "
-        "The reference clip and its transcript only describe the voice. "
-        "**Text to generate** is the line that gets spoken. "
-        "Write ordinary text. You do not need `S0:` or a speaker token. "
-        "This page is a single-speaker preview. Multi-speaker scenes, streaming, "
-        f"and the HTTP API are in the [SGLang Omni cookbook]({COOKBOOK_URL}).",
-        elem_id="fa-intro",
-    )
-    gr.HTML(
-        f"""
-        <nav id="fa-links">
-          <a href="{COOKBOOK_URL}" target="_blank" rel="noreferrer">SGLang Omni cookbook</a>
-          <a href="{SGLANG_OMNI_URL}" target="_blank" rel="noreferrer">⭐ Star SGLang Omni</a>
-          <a href="{TECHNICAL_REPORT_URL}" target="_blank" rel="noreferrer">Technical report</a>
-          <a href="{FISH_AUDIO_URL}" target="_blank" rel="noreferrer">Fish Audio</a>
-          <a href="{FISH_SPEECH_URL}" target="_blank" rel="noreferrer">Fish Speech</a>
-          <a href="{MODEL_URL}" target="_blank" rel="noreferrer">Model card</a>
-        </nav>
-        """
-    )
-    consent_ok = gr.State(False)
-    with gr.Row():
-        with gr.Column(scale=3, elem_id="fa-form"):
-            with gr.Accordion("🎤 Voice consent gate · required before cloning", open=True):
-                gr.Markdown(
-                    "A voice is cloned only when its speaker explicitly agrees "
-                    "([voice consent gate](https://huggingface.co/blog/voice-consent-gate)): "
-                    "generate a unique sentence, read it aloud into your microphone, "
-                    "and we verify it was said. The same recording then serves as the "
-                    "reference audio — spoken consent and cloning input in one."
-                )
-                consent_sentence_btn = gr.Button("1 · Generate a consent sentence", size="sm")
-                consent_sentence = gr.Textbox(
-                    label="Read this sentence aloud",
-                    interactive=False,
-                    lines=2,
-                )
-                consent_recording = gr.Audio(
-                    label="2 · Record yourself (microphone only)",
-                    sources=["microphone"],
-                    type="filepath",
-                )
-                consent_verify_btn = gr.Button("3 · Verify consent", size="sm", variant="secondary")
-                consent_status = gr.Markdown("")
-            reference_audio = gr.Audio(
-                label="Reference audio",
-                sources=["upload", "microphone"],
-                type="filepath",
-            )
-            gr.Markdown(
-                "The voice to copy. Upload a file or record with the microphone. "
-                "Use one speaker, between **0.5 and 30 seconds**. "
-                "A longer clip is rejected. It is not trimmed, so the transcript "
-                "still matches the whole recording."
-            )
-            reference_text = gr.Textbox(
-                label="Reference transcript",
-                placeholder="Hello, nice to meet you. My name is Chenyang.",
-                info="Exact words spoken in the reference audio. This is not the line that will be generated.",
-                lines=2,
-            )
-            text_to_generate = gr.Textbox(
-                label="Text to generate",
-                placeholder="Have a nice day and enjoy Southern California sunshine.",
-                info="The new line this voice should say. Plain text is enough.",
-                lines=3,
-            )
-            generate_button = gr.Button("Generate", variant="primary")
-            generated_audio = gr.Audio(label="Generated audio", interactive=False)
-            gr.Markdown(
-                "One playable clip of **Text to generate** in the reference voice. "
-                "It should not repeat the reference recording."
-            )
-        with gr.Column(scale=2, elem_id="fa-guide"):
-            gr.Markdown(
-                "### What each field means\n"
-                "**Reference audio** is the voice sample. "
-                "It is a condition, not a previous take.\n\n"
-                "**Reference transcript** is what that sample actually says. "
-                "The model uses it to line the voice up with the words. "
-                "Those words are not synthesized again.\n\n"
-                "**Text to generate** is the only text spoken in the output. "
-                "It is assigned to this one voice automatically.\n\n"
-                "**Generate** runs one synthesis. "
-                "There is no second revision pass.\n\n"
-                "**Generated audio** is the result. Play it here or download it.\n\n"
-                "### Example\n"
-                "Reference transcript: "
-                "`Hello, nice to meet you. My name is Chenyang.`\n\n"
-                "Text to generate: "
-                "`Have a nice day and enjoy Southern California sunshine.`\n\n"
-                "The clip should say the second sentence, not the introduction.\n\n"
-                "### Run the full model on SGLang Omni\n"
-                f"- [Fish Audio Drama 3 cookbook]({COOKBOOK_URL}): "
-                "install, voice cloning, and the HTTP API\n"
-                f"- [SGLang Omni]({SGLANG_OMNI_URL}): "
-                "the GitHub repo — please give it a star\n"
-                f"- [Technical report]({TECHNICAL_REPORT_URL}): "
-                "architecture of the 35B MoE model\n"
-                f"- [Fish Audio Playground]({FISH_AUDIO_URL}): "
-                "try Fish Audio voices in the browser\n"
-                f"- [Fish Speech]({FISH_SPEECH_URL}): "
-                "source repository for the model\n"
-                f"- [drama-3 model card]({MODEL_URL}): "
-                "weights, license, and checkpoint files"
-            )
-    gr.Markdown(
-        "### Acknowledgments\n"
-        "It took us over half a year to bring SGLang Omni back to life. "
-        "Thanks to all the people who helped us along the way.\n\n"
-        f"Thanks so much for [AK]({AK_URL}) help to set up the space.\n\n"
-        "**SGLang Omni team:** "
-        f"[Chenyang Zhao]({CHENYANG_URL}), Jingwen Gu, Shenggui Li, Richard Wang, "
-        "Yichi Zhang, Xuesong Ye, Jiaxin Deng, Ratish P, Jiaxuan Luo, Xinyu Lu, "
-        "Wenyao Gao, Ao Sun, Jintao Quaige Li, Yuhao Chen, Yuhan Zhou, "
-        "Xuxiang Tian, Ruoyi Zhang, Junnan Li, Yixiao Zeng, Chenchen Hong, "
-        "Huapeng Zhou, Yang Gao, Chaoye Fei, Yiyang Zh, Kang, Yitian Gong, "
-        "Yuan Luo, Hao Jin, Yifei Gao, Gaokai Zhang, Lifan Shen, Hui Shen, "
-        "Kejing Xia, Jinxin Li, Dongming Shen, Jun Liu, Yue Yin, Zhihao Guo, "
-        "Jisen Li, Ruoyu He, Xuanyi Li, Chenye Zhu, Liang Zhen, Xiaodong Ye, "
-        "Yizhi Zhen, Kaizuo Li, Chonghan Liu, Wei Guan, Zhuofeng Li\n\n"
-        "**FishAudio Team:** Shijia Liao, Shidong Li, Xuehao Yang, Jiahua Liu",
-        elem_id="fa-ack",
-    )
-    consent_sentence_btn.click(
-        fn=on_generate_sentence,
-        outputs=[consent_sentence, consent_ok, consent_status],
-        api_name=False,
-    )
-    consent_verify_btn.click(
-        fn=on_verify,
-        inputs=[consent_recording, consent_sentence],
-        outputs=[reference_audio, reference_text, consent_ok, consent_status],
-        api_name=False,
-    )
-    generate_button.click(
-        fn=synthesize,
-        inputs=[reference_audio, reference_text, text_to_generate, consent_ok],
-        outputs=generated_audio,
-        api_name="synthesize",
-        concurrency_limit=1,
-    )
-
-demo.queue(max_size=16, default_concurrency_limit=1)
-
-if __name__ == "__main__":
-    demo.launch(css=FISH_CSS + WORKFLOW_CSS)

drama3_space_ext.py

@@ -0,0 +1,292 @@
+"""Drama 3 prompt layouts that SGLang Omni's /v1/audio/speech cannot express.
+
+The stock Drama 3 preprocessor builds one layout: optional voice references in the
+system turn, then a single user turn (Fish Audio Drama 3 Technical Report, Figure 11).
+Drama 3 was also trained on two layouts this Space uses:
+
+- casting a voice from a text description, with no reference audio (Figure 13), and
+- multi-turn directing: a note in a later user turn revises the take the model just
+  produced, which sits in the prompt as an earlier assistant turn (Figure 12).
+
+Drama3SpacePreprocessor reads a JSON envelope from the request text and builds those
+layouts. Requests without the envelope take the stock path unchanged. The engine
+listens on localhost only, so the Space's web server is the only client that can send
+an envelope.
+
+SGLang Omni declares each stage's factory in the model's config class and refuses to
+override it from the command line. The Dockerfile instead appends one line to
+sglang_omni/models/fishaudio_drama3/stages.py that rebinds the name its preprocessing
+factory builds, Drama3SpeechPreprocessor, to this class.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import json
+import logging
+from collections.abc import Mapping, Sequence
+
+import torch
+
+from sglang_omni.models.fishaudio_drama3.prompt import (
+    GENERATION_START,
+    MESSAGE_END,
+    REFERENCE_INSTRUCTION,
+    SEMANTIC_BEGIN_TOKEN,
+    SPEAKER_TAG_PATTERN,
+    TTS_INSTRUCTION,
+    UNICODE_WHITE_SPACE,
+    validate_caller_text,
+    with_speaker_tag,
+)
+from sglang_omni.models.fishaudio_drama3.request_builders import REQUEST_ERROR_PREFIX
+from sglang_omni.models.fishaudio_drama3.speech_request import (
+    Drama3SpeechPreprocessor,
+    speech_inputs,
+)
+from sglang_omni.proto.request import OmniRequest, StagePayload
+
+logger = logging.getLogger(__name__)
+
+ENVELOPE_MARKER = "@@drama3-space@@"
+# The scene-dubbing system instruction of Figure 13. The engine samples only audio
+# and end-of-speech tokens, so a request carries one line and never needs the
+# separator; the instruction keeps the layout the model was trained on.
+SCENE_INSTRUCTION = (
+    "Voice every line of the script in order, with <|audio_separator|> "
+    "between consecutive lines."
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class Character:
+    name: str
+    description: str
+    scene: str | None
+
+
+@dataclasses.dataclass(frozen=True)
+class Envelope:
+    """One request: the line to speak now and how its prompt is laid out.
+
+    The request's references hold voice_references voice clips first, then one
+    clip per history turn, in order.
+    """
+
+    text: str
+    voice_references: int
+    history: tuple[str, ...]
+    character: Character | None
+
+
+def request_error(message: str) -> ValueError:
+    return ValueError(f"{REQUEST_ERROR_PREFIX} {message}")
+
+
+def single_line(value: object, field_name: str, *, optional: bool = False) -> str | None:
+    if value is None and optional:
+        return None
+    if not isinstance(value, str) or not value.strip(UNICODE_WHITE_SPACE):
+        raise request_error(f"{field_name} must be a non-empty string")
+    if "\n" in value or "\r" in value:
+        raise request_error(f"{field_name} must be a single line")
+    validate_caller_text(value, field_name)
+    return value.strip(UNICODE_WHITE_SPACE)
+
+
+def parse_envelope(body: str, reference_count: int) -> Envelope:
+    try:
+        spec = json.loads(body)
+    except json.JSONDecodeError as error:
+        raise request_error(f"malformed envelope: {error}") from error
+    if not isinstance(spec, Mapping):
+        raise request_error("the envelope must be a JSON object")
+
+    text = spec.get("text")
+    if not isinstance(text, str):
+        raise request_error("envelope text must be a string")
+    validate_caller_text(text, "text")
+
+    history = spec.get("history") or []
+    if not isinstance(history, list) or not all(isinstance(t, str) for t in history):
+        raise request_error("envelope history must be a list of strings")
+    for index, turn_text in enumerate(history):
+        validate_caller_text(turn_text, f"history.{index}")
+
+    voice_references = spec.get("voice_references", 0)
+    if (
+        isinstance(voice_references, bool)
+        or not isinstance(voice_references, int)
+        or voice_references < 0
+        or voice_references + len(history) != reference_count
+    ):
+        raise request_error(
+            f"voice_references {voice_references!r} and {len(history)} history "
+            f"turns do not account for {reference_count} references"
+        )
+
+    character = None
+    raw_character = spec.get("character")
+    if raw_character is not None:
+        if not isinstance(raw_character, Mapping):
+            raise request_error("envelope character must be an object")
+        if voice_references:
+            raise request_error("a described character cannot also have voice references")
+        name = single_line(raw_character.get("name"), "character.name")
+        if ":" in name:
+            raise request_error("character.name must not contain ':'")
+        character = Character(
+            name=name,
+            description=single_line(
+                raw_character.get("description"), "character.description"
+            ),
+            scene=single_line(raw_character.get("scene"), "character.scene", optional=True),
+        )
+
+    return Envelope(
+        text=text,
+        voice_references=voice_references,
+        history=tuple(history),
+        character=character,
+    )
+
+
+def reference_text_block(transcript: str) -> str:
+    """The stock reference transcript: the speaker tag follows "Text: " directly."""
+    tagged = with_speaker_tag(transcript, 0)
+    body = tagged.lstrip(UNICODE_WHITE_SPACE)
+    tag = SPEAKER_TAG_PATTERN.match(body)
+    assert tag is not None
+    leading = tagged[: len(tagged) - len(body)]
+    return tag.group(0) + leading + body[tag.end() :]
+
+
+class Drama3SpacePreprocessor(Drama3SpeechPreprocessor):
+    """The stock preprocessor, plus described characters and directing history."""
+
+    def __init__(self, **kwargs: object) -> None:
+        super().__init__(**kwargs)
+        # server.py looks for this line to confirm the stages.py hook took effect.
+        logger.warning("Drama 3 Space preprocessor active")
+
+    def __call__(self, payload: StagePayload) -> StagePayload:
+        text, descriptors = speech_inputs(payload.request.inputs)
+        if not text.startswith(ENVELOPE_MARKER):
+            return super().__call__(payload)
+        envelope = parse_envelope(text[len(ENVELOPE_MARKER) :], len(descriptors))
+
+        # The stock path checks the sampling fields and the line itself; the rows it
+        # builds for that reference-free probe are replaced below.
+        probe = StagePayload(
+            request_id=payload.request_id,
+            request=OmniRequest(
+                inputs=self.turn_text(envelope.text, envelope.character),
+                params=payload.request.params,
+                metadata=payload.request.metadata,
+            ),
+            data=None,
+        )
+        checked = super().__call__(probe)
+
+        codes = [
+            self.reference_audio_codes(descriptor, index)
+            for index, descriptor in enumerate(descriptors)
+        ]
+        voice_codes = codes[: envelope.voice_references]
+        history_codes = codes[envelope.voice_references :]
+        transcripts = []
+        for index, descriptor in enumerate(descriptors[: envelope.voice_references]):
+            transcript = descriptor.get("text")
+            if not isinstance(transcript, str):
+                raise request_error(f"references.{index}.text must be a string")
+            validate_caller_text(transcript, f"references.{index}")
+            transcripts.append(transcript)
+
+        segments: list[str | torch.Tensor] = ["<|im_start|>system\n"]
+        if envelope.character is not None:
+            character = envelope.character
+            lines = [f"Scene: {character.scene}"] if character.scene else []
+            lines += [
+                "Characters:",
+                f"- {character.name}: {character.description}",
+                SCENE_INSTRUCTION,
+            ]
+            segments.append("\n".join(lines))
+        else:
+            segments.append(REFERENCE_INSTRUCTION if voice_codes else TTS_INSTRUCTION)
+            for index, (transcript, reference) in enumerate(
+                zip(transcripts, voice_codes, strict=True), 1
+            ):
+                segments += [
+                    f"\n\n[Reference {index}]\nText: "
+                    f"{reference_text_block(transcript)}\nSpeech: ",
+                    reference,
+                ]
+        segments.append(MESSAGE_END)
+        for turn_text, take in zip(envelope.history, history_codes, strict=True):
+            segments += [
+                "<|im_start|>user\n",
+                self.turn_text(turn_text, envelope.character),
+                MESSAGE_END,
+                GENERATION_START,
+                take,
+                MESSAGE_END,
+            ]
+        segments += [
+            "<|im_start|>user\n",
+            self.turn_text(envelope.text, envelope.character),
+            MESSAGE_END,
+            GENERATION_START,
+        ]
+
+        rows = self.segment_rows(segments)
+        if rows.shape[0] > self.max_prompt_rows:
+            raise request_error(
+                f"a prompt of {rows.shape[0]} rows exceeds the maximum allowed "
+                f"length of {self.max_prompt_rows} rows"
+            )
+        return StagePayload(
+            request_id=payload.request_id,
+            request=payload.request,
+            data=dataclasses.replace(checked.data, prompt_rows=rows),
+        )
+
+    @staticmethod
+    def turn_text(text: str, character: Character | None) -> str:
+        if character is not None:
+            return f"{character.name}: {text}"
+        return with_speaker_tag(text, 0)
+
+    def segment_rows(self, segments: Sequence[str | torch.Tensor]) -> torch.Tensor:
+        """Text rows are [token, 0, ...]; audio rows [c0 + semantic_begin, c0..cN-1].
+
+        Every literal is tokenized on its own, as in the stock builder: tokenizing
+        the joined string changes ids at the seams.
+        """
+        semantic_begin_id = self.tokenizer.token_to_id(SEMANTIC_BEGIN_TOKEN)
+        rows: list[torch.Tensor] = []
+        for segment in segments:
+            if isinstance(segment, str):
+                token_ids = torch.tensor(
+                    self.tokenizer.encode(segment).ids, dtype=torch.long
+                )
+                text_rows = torch.zeros(
+                    (token_ids.numel(), self.num_codebooks + 1), dtype=torch.long
+                )
+                text_rows[:, 0] = token_ids
+                rows.append(text_rows)
+            else:
+                audio_codes = segment.to(device="cpu", dtype=torch.long)
+                if (
+                    audio_codes.ndim != 2
+                    or audio_codes.shape[0] != self.num_codebooks
+                    or audio_codes.shape[1] == 0
+                ):
+                    raise request_error(
+                        f"audio codes must be [{self.num_codebooks}, frames]; got "
+                        f"{tuple(audio_codes.shape)}"
+                    )
+                rows.append(
+                    torch.cat([audio_codes[:1] + semantic_begin_id, audio_codes]).T
+                )
+        return torch.cat(rows)

endpoint/.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+status.json
+logs.txt
+verification.json
+*.wav
+comparisons/

endpoint/PERFORMANCE.md

@@ -0,0 +1,78 @@
+# H200 performance investigation — 14 October 2026
+
+These initial measurements precede the later transport optimizations described
+in [README.md](README.md#measure-ui-and-network-overhead), including the current
+2.5-second opening threshold, streaming compression, and HTTP/2 connection reuse.
+
+The H200 was correctly detected and already using CUDA graphs, Triton sampling,
+and fused tower layers. One-step asynchronous decoding was disabled. It is now
+enabled for batch sizes starting at one, on the same endpoint, image, model
+revision, GCP region, replica count, and idle scaling policy.
+
+## Warm endpoint measurements
+
+Each run includes four measured Surprise and four measured Clone requests,
+plus one excluded warm-up per mode. The same 31-word script, exact reference
+PCM, and temperature 1 / top-p 0.9 / top-k 4097 are used throughout. Audio lengths
+vary because this model has no request seed. Times are measured from the client
+in Europe and include authentication, network, and reference uploads.
+
+| Measurement (median) | Before | Async decode | Async + lossless FLAC reference |
+| --- | ---: | ---: | ---: |
+| Streaming audio seconds per elapsed second, after first chunk | 8.24× | 10.45× | 10.13× |
+| First audio, Surprise | 775 ms | 762 ms | 760 ms |
+| First audio, Clone | 1,573 ms | 1,450 ms | 1,272 ms |
+| Normal stops | 8/8 | 8/8 | 8/8 |
+
+Async decoding improved measured streaming speed by about 27%. Separate steady
+single-request decode log samples increased from a median 168 to 197 frames/s
+(about 17%). These measurements describe different layers: the streaming ratio
+also includes vocoding, chunk timing, and transfer; it is not a pure GPU benchmark.
+
+The cloned reference shrank from 862,106 to 417,548 bytes with FLAC, with identical
+decoded PCM16 samples. The Space now caches that compressed representation for
+voice uploads and uses FLAC for directing-history audio too. Browser downloads
+remain WAV. Native API callers can also send `media_type="audio/flac"`.
+
+## Why the Space can feel slower
+
+The current Space runs on CPU and proxies the same H200 endpoint. It is no longer
+an independent GPU control. No controlled A6000-versus-H200 conclusion can be
+drawn from these tests.
+
+Before the packet change, a valid three-second opening was sent as one large
+base64 SSE event. The client cannot decode that event until the whole line has
+arrived. In the async test, the Space measured about 0.8–1.1 seconds to prepare
+first audio, but the client's complete audio event arrived several seconds later.
+This is evidence of delivery overhead, not slow GPU decoding.
+
+The Space now splits validated audio into packets of at most half a second.
+It keeps the same three-second opening validation, silence trimming, early retry,
+quota accounting, and cancellation behavior. First-audio timings from the server
+describe audio preparation; browser playback also includes delivery and rendering.
+
+GCP `us-south1` adds a transatlantic trip for European clients. The endpoint also
+still scales to zero after 60 idle minutes; a cold start has to reload weights and
+capture GPU graphs. The observed configuration restart took about 3 minutes
+20 seconds to become healthy, and the first speech request then took about six
+seconds. Neither time is included in the warm medians above.
+
+## Validation and limits
+
+50 local checks cover streaming packets, exact reference PCM preservation, Gradio,
+Workflow, quotas, cancellation, failed-opening retries, and safe serving updates.
+Live testing also ran four simultaneous requests, cancelled one after first audio,
+and verified that the other three and a subsequent request stopped normally.
+
+Four endpoint samples per before/async run were transcribed through the same
+Whisper service. Before: four exact scripts. After: three exact scripts and one
+with the final word missing from transcription and a long pause. All requests
+stopping normally does not prove that every script is complete. These small samples
+are a speed investigation, not evidence that silence or completion issues are fixed.
+The requested sampling distribution is unchanged by the serving optimization.
+
+Reproduce using `benchmark_speed.py`; raw audio, arrival times, and summaries live
+under the Git-ignored `endpoint/comparisons/speed-*` directories. Published upstream
+[H200 benchmarks](https://sgl-project.github.io/sglang-omni/cookbook/fishaudio_drama3.html#benchmarks)
+use top-k 30 and pre-encoded voice references, so their sub-100 ms first-audio
+results should not be compared directly with this authenticated remote API test.

endpoint/README.md

@@ -0,0 +1,143 @@
+# Drama 3 on a GCP H200
+
+The inference server is deployed as **HuggingFaceM4/drama-3-h200** on Hugging Face
+Inference Endpoints, backed by **one Nvidia H200 (141 GB)** in GCP **us-south1**.
+At deployment, the hardware catalog quotes **$5/hour while running**. There is at
+most one replica, and it scales to zero after **60 minutes without requests**.
+The endpoint requires a Hugging Face token with access to the organization.
+
+## Server and reproducibility
+
+`deployment.json` pins both the `fishaudio/drama-3` model revision and the existing
+Drama-3 Space image by digest. The image includes SGLang Omni at
+`22d4ace88f9547777b819a69d7691204bb9ce847` and the voice-description / directing-history
+prompt extension. The endpoint runs SGLang Omni directly on port 8000 with model
+weights mounted at `/repository`; it exposes `/health` and `/v1/audio/speech`.
+H200 enables the engine's Hopper-compatible sampling and fused tower kernels.
+Asynchronous decoding overlaps the next GPU step with CPU result processing,
+including for one active request (`async_decode_min_batch_size=1`).
+The server caps each generation at 1200 audio frames (about 56 seconds).
+
+The Space now runs on CPU and sends generation to this endpoint using `DRAMA_ENDPOINT_URL` and a server-only `DRAMA_ENDPOINT_TOKEN` secret. Its CPU `Dockerfile` is separate from the retained `Dockerfile.engine` GPU recipe; editing the UI does not rebuild the endpoint.
+
+The endpoint serves the native speech API. The Space's silence trimming, early
+retry, take storage, presets, quotas, and Gradio controls belong to its web server.
+
+To recreate the endpoint, using an authenticated Hugging Face environment:
+
+```sh
+python endpoint/deploy.py
+python endpoint/watch.py
+python endpoint/verify.py
+```
+
+To apply updated serving arguments to this existing, matching deployment:
+
+```sh
+python endpoint/deploy.py --update-serving
+python endpoint/watch.py
+```
+
+This restarts the server. It checks the pinned image, model revision, command,
+provider, and hardware first; it changes only container arguments.
+
+The deploy script checks quota and reuses an existing deployment with the same
+model and hardware instead of creating another server. Runtime status, logs,
+verification results, and sample WAVs are ignored by Git. `verify.py` checks a
+plain script, an expressive script, and voice cloning through the streaming API.
+
+## Call the server
+
+The endpoint URL is discovered from the API rather than hard-coded:
+
+```python
+import os
+import httpx
+from huggingface_hub import HfApi
+
+endpoint = HfApi().get_inference_endpoint(
+    "drama-3-h200", namespace="HuggingFaceM4"
+)
+response = httpx.post(
+    endpoint.url + "/v1/audio/speech",
+    headers={"Authorization": "Bearer " + os.environ["HF_TOKEN"]},
+    json={
+        "input": "Good evening. Welcome to Drama Three.",
+        "response_format": "wav",
+        "max_new_tokens": 480,
+        "temperature": 1.0,
+        "top_p": 0.9,
+        "top_k": 4097,
+    },
+    timeout=300,
+)
+response.raise_for_status()
+with open("speech.wav", "wb") as audio:
+    audio.write(response.content)
+```
+
+For streaming, set `stream=true`, `response_format="pcm"`, and
+`stream_format="sse"`. Audio arrives in `speech.audio.delta` events as base64 PCM16;
+`speech.audio.done` carries the finish reason. Close the response after that event.
+For cloning, send `references` containing base64 audio, its media type, and transcript.
+Lossless FLAC (`audio/flac`) reduces reference uploads while preserving PCM samples;
+the Space uses this format for clones and directing history.
+A scaled-to-zero endpoint needs time to reload the model on the next request.
+Native top-k must be a positive integer: 4097 includes all 4096 codes plus EOS.
+Omitting it uses upstream's capped default of 30; zero is not accepted by this API.
+
+## Measure streaming speed
+
+```sh
+python endpoint/benchmark_speed.py --out endpoint/comparisons/speed --repeats 4 --include-space
+```
+
+The benchmark keeps temperature 1, top-p 0.9, and uncapped top-k. It saves audio,
+chunk arrival times, first-audio latency, and streaming speed for repeated Surprise
+and Clone takes. One warm-up request per voice mode is excluded from medians.
+Client timings include network and authentication; streaming speed excludes the
+first chunk but still includes vocoding and transport. Audio is stochastic and
+the model has no request seed. The current CPU Space uses this same endpoint,
+so this measures UI/proxy overhead rather than comparing two GPUs.
+Use `--reference-format flac` to measure lossless compressed reference uploads.
+See [PERFORMANCE.md](PERFORMANCE.md) for the measured before/after results and limits.
+
+## Measure UI and network overhead
+
+```sh
+python endpoint/benchmark_transport.py --out endpoint/comparisons/transport --repeats 3
+```
+
+This uses HTTP/2 and measures the native endpoint, Space `/api/generate`, and the
+actual Gradio stage queue with the same script and sampling settings. It records
+first audio, completion, HTTP version, compressed wire bytes, and server timing
+breakdowns. One warm-up per route is excluded. Use `--compression identity` to
+disable response compression for a control run.
+
+The CPU Space maintains an HTTP/2 client with a 120-second idle connection lifetime.
+Successful endpoint health or speech responses avoid another health request for
+30 seconds; there are no background pings, so idle scale-to-zero still works.
+The Space compresses SSE messages with gzip and flushes every packet immediately.
+This covers both the proxy and Gradio's event stream, without compressing internal
+Python iterators or changing the PCM audio, voice references, or sampling settings.
+Clients that do not accept gzip continue receiving plain SSE. Opening validation
+uses 2.5 seconds of retained audio. The early engine windows total about 2.97
+seconds; the previous 3-second threshold could miss that boundary and wait for
+another 64-frame window. Failed openings below the threshold are still withheld
+and retried; their trailing silence cannot push them into playback.
+
+## Compare with the Space
+
+```sh
+python endpoint/compare.py --out endpoint/comparisons/run --repeats 4
+python endpoint/analyze_comparison.py --out endpoint/comparisons/run
+```
+
+These scripts record 24 requests per deployment across six English scripts, with
+matching model revision, prompts, references, sampling controls, and text-dependent
+generation limits. The Space includes its production trimming and retry behavior;
+the endpoint returns native audio. This compares the deployed services, rather
+than isolating GPU hardware. Audio, request outcomes, transcriptions from the
+same Whisper service, and a paired audio report are saved under the output folder.
+The report measures script completion and quiet gaps, with explicit ASR limits;
+it does not rate subjective timbre or expressiveness.

endpoint/analyze_comparison.py

@@ -0,0 +1,155 @@
+"""Transcribe complete generated clips through the same Whisper service; score both."""
+import argparse
+from collections import Counter
+import html
+import io
+import json
+from pathlib import Path
+import re
+import statistics
+
+import httpx
+import numpy as np
+import soundfile as sf
+from huggingface_hub import get_token
+
+SPACE = 'https://andito-drama-3.hf.space'
+
+
+def words(text):
+    return re.findall(r"[a-z0-9]+(?:'[a-z0-9]+)?", text.lower().replace('’', "'"))
+
+
+def alignment(reference, transcript):
+    ref, hyp = words(reference), words(transcript)
+    dp = [[0] * (len(hyp)+1) for _ in range(len(ref)+1)]
+    for i in range(len(ref)+1):
+        dp[i][0] = i
+    for j in range(len(hyp)+1):
+        dp[0][j] = j
+    for i in range(1, len(ref)+1):
+        for j in range(1, len(hyp)+1):
+            dp[i][j] = min(dp[i-1][j]+1, dp[i][j-1]+1, dp[i-1][j-1]+(ref[i-1] != hyp[j-1]))
+    i, j, counts, missing, extra = len(ref), len(hyp), Counter(), [], []
+    while i or j:
+        if i and j and dp[i][j] == dp[i-1][j-1] + (ref[i-1] != hyp[j-1]):
+            if ref[i-1] != hyp[j-1]:
+                counts['substitutions'] += 1
+                missing.append(ref[i-1]); extra.append(hyp[j-1])
+            i -= 1; j -= 1
+        elif i and dp[i][j] == dp[i-1][j] + 1:
+            counts['deletions'] += 1; missing.append(ref[i-1]); i -= 1
+        else:
+            counts['insertions'] += 1; extra.append(hyp[j-1]); j -= 1
+    return {'word_error_rate': dp[-1][-1] / max(1, len(ref)),
+            'word_coverage': (len(ref)-counts['deletions']-counts['substitutions']) / max(1, len(ref)),
+            'extra_word_fraction': counts['insertions'] / max(1, len(ref)),
+            'missing_or_changed_words': list(reversed(missing)), 'extra_or_changed_words': list(reversed(extra)),
+            **{k: counts[k] for k in ['substitutions','deletions','insertions']}}
+
+
+def transcribe(client, wav_path):
+    samples, rate = sf.read(wav_path, dtype='float32')
+    # Split only when necessary: the Space's reference API caps uploads at 20 s.
+    # No overlap, so transcripts cannot acquire artificial repetitions.
+    parts = []
+    details = []
+    for index, start in enumerate(range(0, len(samples), 18*rate)):
+        chunk = samples[start:start+18*rate]
+        if len(chunk) == 0 or np.sqrt(np.mean(chunk**2)) < .00003:
+            details.append({'part': index, 'note': 'Effectively silent segment'})
+            continue
+        buffer = io.BytesIO()
+        sf.write(buffer, chunk, rate, format='WAV', subtype='PCM_16')
+        response = client.post(SPACE + '/api/voices', files={'audio': (f'comparison-{index}.wav', buffer.getvalue(), 'audio/wav')})
+        if response.status_code == 400 and 'too short' in response.text:
+            details.append({'part': index, 'note': 'Speech segment too short for reference transcription'})
+            continue
+        response.raise_for_status()
+        payload = response.json()
+        if not payload.get('transcribed'):
+            raise RuntimeError('Whisper transcription was unavailable; cannot compare word accuracy')
+        parts.append(payload['transcript'])
+        details.append({'part': index, 'transcript': payload['transcript'], 'language': payload.get('language')})
+    return ' '.join(parts).strip(), details
+
+
+def summarize(results):
+    summary = {}
+    for source in ['space','endpoint']:
+        rows = [r for r in results if r['source'] == source]
+        normal = [r for r in rows if r.get('finish_reason') == 'stop' and not r.get('error')]
+        complete = [r for r in rows if r.get('script_complete')]
+        clean = [r for r in rows if r.get('clean_take')]
+        retry_successes = [r for r in clean if r['retry_count']]
+        summary[source] = {'takes': len(rows), 'normal_termination': len(normal), 'script_complete': len(complete),
+                           'clean_takes': len(clean), 'long_internal_pauses': sum(r['longest_internal_quiet_s'] > 2 for r in rows),
+                           'long_trailing_silence': sum(r['trailing_quiet_s'] > 2 for r in rows),
+                           'early_errors': sum(bool(r.get('error')) for r in rows),
+                           'transcription_incomplete': sum(bool(r.get('transcription_incomplete')) for r in rows),
+                           'stalled_or_length_capped': sum(r.get('finish_reason') in {'stall','length','timeout','incomplete'} for r in rows),
+                           'retried_takes': sum(bool(r['retry_count']) for r in rows), 'successful_retries': len(retry_successes),
+                           'clean_without_retry': len(clean)-len(retry_successes),
+                           'max_internal_quiet_s': max(r['longest_internal_quiet_s'] for r in rows),
+                           'max_trailing_quiet_s': max(r['trailing_quiet_s'] for r in rows),
+                           'median_word_error_rate': statistics.median(r['word_error_rate'] for r in rows),
+                           'median_first_audio_client_ms': statistics.median(r['first_audio_client_ms'] for r in rows if r['first_audio_client_ms'] is not None)}
+    return summary
+
+
+def write_report(out, results, summary):
+    escape = html.escape
+    cards = []
+    keys = list(dict.fromkeys((r['case'], r['repeat']) for r in results))
+    for case, repeat in keys:
+        pair = [r for r in results if r['case'] == case and r['repeat'] == repeat]
+        panels = []
+        for source in ['space','endpoint']:
+            r = next(x for x in pair if x['source'] == source)
+            player = f'<audio controls preload="none" src="{escape(r["audio_file"])}"></audio>' if r.get('audio_file') else '<p>No playable take</p>'
+            details = (f"Finished: {r.get('finish_reason', 'error')} · Script coverage: {r['word_coverage']:.0%} · "
+                       f"Internal quiet: {r['longest_internal_quiet_s']:.2f} s · Trailing quiet: {r['trailing_quiet_s']:.2f} s · Retries: {r['retry_count']}")
+            panels.append(f'<section><h3>{escape(source.title())}</h3>{player}<p class="metrics">{escape(details)}</p><p>{escape(r.get("transcript", ""))}</p><p class="failure">{escape(r.get("error", ""))}</p></section>')
+        cards.append(f'<article><h2>{escape(case)} · take {repeat}</h2><p class="script">{escape(pair[0]["expected"])}</p><div class="pair">{"".join(panels)}</div></article>')
+    rows = []
+    for label, key in [('Completed script', 'script_complete'), ('Completed without long silence', 'clean_takes'), ('Long internal pauses (>2 s)', 'long_internal_pauses'), ('Long trailing silence (>2 s)', 'long_trailing_silence'), ('Stalled / reached generation limit', 'stalled_or_length_capped'), ('Error before playback', 'early_errors'), ('Automatic retries', 'retried_takes')]:
+        rows.append(f'<tr><th>{escape(label)}</th><td>{summary["space"][key]}</td><td>{summary["endpoint"][key]}</td></tr>')
+    page = f'''<!doctype html><html lang="en"><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1"><title>Drama 3 — Space vs H200</title><style>body{{font:16px system-ui;max-width:1100px;margin:auto;padding:32px;background:#f6f4ee;color:#24211e}}h1{{font-size:32px}}article{{background:white;border:1px solid #d9d5cd;border-radius:16px;padding:22px;margin:22px 0}}.pair{{display:grid;grid-template-columns:1fr 1fr;gap:24px}}audio{{width:100%}}.metrics{{font-size:13px;color:#665f56}}.script{{font-weight:600}}.failure{{color:#a53124}}table{{border-collapse:collapse;background:white;width:100%}}th,td{{padding:10px;text-align:left;border-bottom:1px solid #ddd}}.method{{line-height:1.5}}@media(max-width:700px){{.pair{{grid-template-columns:1fr}}}}</style><h1>Drama 3: Space vs H200 endpoint</h1><p class="method">24 requests per server: six scripts, four repetitions. Matching model revision, prompts, reference audio, sampling settings, and generation limits. Space: RTX PRO 6000 with silence trimming and up to one retry. Endpoint: H200 with native, unprocessed audio. Random draws differ, and Surprise can choose different voices. This compares the deployed services; it does not isolate hardware alone.</p><p class="method">Same Whisper service transcribes both outputs. Clips longer than 18 seconds are transcribed in non-overlapping parts. Script completion means at least 90% of words matched and no substantial added-word sequence (insertions ≤25% of script length). Clean takes also require a normal stop and no internal or trailing quiet segment over 2 seconds. Quiet uses the Space's RMS threshold of 0.003; very soft speech may count as quiet. Segments with less than one second of speech after trimming cannot be transcribed by the reference API; they are marked unverified and do not count as a verified completed script. ASR can mishear words, especially whispers, laughter and split boundaries. These are reliability measures, not a subjective voice-quality score.</p><table><tr><th>Result (out of 24)</th><th>Space</th><th>H200 endpoint</th></tr>{''.join(rows)}</table>{''.join(cards)}</html>'''
+    (out / 'comparison.html').write_text(page)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--out', type=Path, required=True)
+    args = parser.parse_args()
+    out = args.out
+    results = json.loads((out / 'generations.json').read_text())
+    analyzed_file = out / 'analyzed.json'
+    existing = json.loads(analyzed_file.read_text()) if analyzed_file.exists() else []
+    cached = {(r['source'], r['case'], r['repeat']): r for r in existing}
+    analyzed = []
+    with httpx.Client(headers={'Authorization': 'Bearer ' + get_token()}, timeout=90, follow_redirects=True) as client:
+        for r in results:
+            key = r['source'], r['case'], r['repeat']
+            if key in cached:
+                analyzed.append(cached[key]); continue
+            if r.get('audio_file'):
+                transcript, details = transcribe(client, out / r['audio_file'])
+            else:
+                transcript, details = '', []
+            r.update(transcript=transcript, transcription_parts=details, **alignment(r['expected'], transcript))
+            r['transcription_incomplete'] = any('too short' in part.get('note', '') for part in details)
+            r['script_complete'] = not r['transcription_incomplete'] and r['word_coverage'] >= .90 and r['extra_word_fraction'] <= .25
+            r['clean_take'] = (r['script_complete'] and r.get('finish_reason') == 'stop' and not r.get('error')
+                               and r['longest_internal_quiet_s'] <= 2 and r['trailing_quiet_s'] <= 2)
+            analyzed.append(r)
+            analyzed_file.write_text(json.dumps(analyzed, indent=2) + '\n')
+            print(json.dumps({k:r.get(k) for k in ['source','case','repeat','transcript','word_coverage','script_complete','clean_take']}), flush=True)
+    summary = summarize(analyzed)
+    (out / 'summary.json').write_text(json.dumps(summary, indent=2) + '\n')
+    write_report(out, analyzed, summary)
+    print('SUMMARY', json.dumps(summary), flush=True)
+
+
+if __name__ == '__main__':
+    main()

endpoint/benchmark_speed.py

@@ -0,0 +1,138 @@
+"""Measure warm streaming speed with unchanged production sampling and references.
+
+Run before and after a serving configuration change. Latencies are client-side;
+stream speed excludes the first chunk, but still includes vocoding and transport.
+The model has no request seed, so compare repeated runs rather than take lengths.
+"""
+import argparse
+import base64
+import json
+from pathlib import Path
+import statistics
+import time
+
+import httpx
+from huggingface_hub import HfApi, get_token
+
+from compare import SPACE, SETTINGS, audio_metrics
+
+TEXT = ("The rain has stopped, and the street is quiet. I left a note on the kitchen table. "
+        "Please read it before you leave, then call me when you reach the station.")
+
+
+def measure(client, url, payload, label, out):
+    started = time.perf_counter()
+    chunks, pcm, finish, server = [], bytearray(), None, {}
+    rate = 44100
+    with client.stream('POST', url, json=payload) as response:
+        response.raise_for_status()
+        headers_s = time.perf_counter() - started
+        rate = int(response.headers.get('x-sample-rate', rate))
+        for line in response.iter_lines():
+            if not line.startswith('data: '):
+                continue
+            event = json.loads(line[6:])
+            kind = event.get('type')
+            if kind == 'start':
+                rate = event['sample_rate']
+                server['start_client_ms'] = round((time.perf_counter() - started) * 1000)
+            if kind in {'speech.audio.delta', 'audio'}:
+                pcm.extend(base64.b64decode(event.get('audio', event.get('pcm', ''))))
+                chunks.append({'at_s': time.perf_counter() - started, 'audio_s': len(pcm) / (rate * 2)})
+            if kind in {'speech.audio.done', 'done'}:
+                finish = event.get('finish_reason')
+                server.update({key: event[key] for key in ['ttfa_ms', 'total_ms', 'queue_ms',
+                                                          'retry_count', 'trimmed_s', 'timings_ms'] if key in event})
+                break
+            if kind == 'error':
+                raise RuntimeError(event.get('message', str(event)))
+        wire_bytes = response.num_bytes_downloaded
+        http_version = response.http_version
+        encoding = response.headers.get('content-encoding', 'identity')
+    wall_s = time.perf_counter() - started
+    stream_speed = None
+    if len(chunks) > 1 and chunks[-1]['at_s'] > chunks[0]['at_s']:
+        stream_speed = ((chunks[-1]['audio_s'] - chunks[0]['audio_s']) /
+                        (chunks[-1]['at_s'] - chunks[0]['at_s']))
+    result = {'label': label, 'headers_ms': round(headers_s * 1000),
+              'first_audio_ms': round(chunks[0]['at_s'] * 1000) if chunks else None,
+              'wall_s': round(wall_s, 3), 'finish_reason': finish,
+              'stream_speed_x': round(stream_speed, 3) if stream_speed else None,
+              'chunks': chunks, **server, **audio_metrics(pcm, rate)}
+    result.update(wire_bytes=wire_bytes, http_version=http_version, content_encoding=encoding)
+    result['overall_speed_x'] = round(result['audio_s'] / wall_s, 3)
+    # Keep raw audio for completion/quality checks, rather than timing alone.
+    import wave
+    with wave.open(str(out / (label + '.wav')), 'wb') as audio:
+        audio.setnchannels(1)
+        audio.setsampwidth(2)
+        audio.setframerate(rate)
+        audio.writeframes(pcm)
+    return result
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--out', type=Path, required=True)
+    parser.add_argument('--repeats', type=int, default=4)
+    parser.add_argument('--include-space', action='store_true')
+    parser.add_argument('--reference-format', choices=['wav', 'flac'], default='wav',
+                        help='FLAC preserves exact PCM samples while reducing reference upload size.')
+    args = parser.parse_args()
+    args.out.mkdir(parents=True, exist_ok=True)
+    endpoint = HfApi().get_inference_endpoint('drama-3-h200', namespace='HuggingFaceM4')
+    if endpoint.status != 'running' or not endpoint.url:
+        raise RuntimeError('Endpoint must be running; cold starts are a separate measurement.')
+    results = []
+    with httpx.Client(headers={'Authorization': 'Bearer ' + get_token()}, timeout=180) as client:
+        status = client.get(SPACE + '/api/status').raise_for_status().json()
+        voice = next(v for v in status['voices'] if v['id'] == 'clone1')
+        reference = client.get(SPACE + '/api/voices/clone1.wav').raise_for_status().content
+        if args.reference_format == 'flac':
+            import io
+            import numpy as np
+            import soundfile as sf
+            samples, rate = sf.read(io.BytesIO(reference), dtype='int16')
+            encoded = io.BytesIO()
+            sf.write(encoded, samples, rate, format='FLAC', subtype='PCM_16')
+            decoded, decoded_rate = sf.read(io.BytesIO(encoded.getvalue()), dtype='int16')
+            if decoded_rate != rate or not np.array_equal(samples, decoded):
+                raise RuntimeError('Reference compression changed the audio samples.')
+            reference = encoded.getvalue()
+        refs = [{'data': base64.b64encode(reference).decode(), 'media_type': 'audio/' + args.reference_format,
+                 'text': voice['transcript']}]
+        metadata = {'endpoint_model_revision': endpoint.raw['model']['revision'],
+                    'container_args': endpoint.raw['model'].get('args'),
+                    'settings': SETTINGS, 'text': TEXT, 'space_gpu': status['gpu'],
+                    'reference_format': args.reference_format, 'reference_bytes': len(reference),
+                    'repeats': args.repeats, 'frame_rate': 44100 / 2048}
+        (args.out / 'metadata.json').write_text(json.dumps(metadata, indent=2) + '\n')
+        for mode in ['surprise', 'clone']:
+            native = {'input': TEXT, 'stream': True, 'response_format': 'pcm', 'stream_format': 'sse',
+                      'max_new_tokens': 500, 'references': refs if mode == 'clone' else None, **SETTINGS}
+            # The first request for each shape is kept but excluded from warm medians.
+            for i in range(args.repeats + 1):
+                targets = ['endpoint', 'space'] if args.include_space else ['endpoint']
+                for target in targets:
+                    payload = native if target == 'endpoint' else {
+                        'text': TEXT, 'voice': {'mode': 'clone', 'id': 'clone1'} if mode == 'clone'
+                        else {'mode': 'random'}, **SETTINGS}
+                    url = endpoint.url + '/v1/audio/speech' if target == 'endpoint' else SPACE + '/api/generate'
+                    result = measure(client, url, payload, f'{target}-{mode}-{i}', args.out)
+                    result.update(target=target, mode=mode, warmup=(i == 0))
+                    results.append(result)
+                    print(json.dumps({k: v for k, v in result.items() if k != 'chunks'}), flush=True)
+                    (args.out / 'results.json').write_text(json.dumps(results, indent=2) + '\n')
+    summary = {}
+    for target in {r['target'] for r in results}:
+        warm = [r for r in results if r['target'] == target and not r['warmup']]
+        summary[target] = {key: round(statistics.median(r[key] for r in warm if r[key] is not None), 3)
+                           for key in ['first_audio_ms', 'wall_s', 'stream_speed_x', 'overall_speed_x']}
+        summary[target]['stopped_normally'] = sum(r['finish_reason'] == 'stop' for r in warm)
+        summary[target]['requests'] = len(warm)
+    (args.out / 'summary.json').write_text(json.dumps(summary, indent=2) + '\n')
+    print(json.dumps(summary, indent=2), flush=True)
+
+
+if __name__ == '__main__':
+    main()

endpoint/benchmark_transport.py

@@ -0,0 +1,135 @@
+"""Compare native endpoint, Space proxy, and the actual Gradio stage event stream."""
+import argparse
+import base64
+import copy
+import json
+from pathlib import Path
+import statistics
+import time
+import uuid
+import wave
+
+import httpx
+from huggingface_hub import HfApi, get_token
+
+from benchmark_speed import TEXT, SETTINGS, measure, SPACE
+
+
+def apply_diff(state, value):
+    if isinstance(value, dict):
+        return value
+    state = copy.deepcopy(state)
+    for operation, path, item in value:
+        if not path:
+            state = item
+            continue
+        parent = state
+        for key in path[:-1]:
+            parent = parent[key]
+        key = path[-1]
+        if operation == 'delete':
+            del parent[key]
+        elif operation == 'append':
+            parent[key] += item
+        elif operation == 'add' and isinstance(parent, list):
+            parent.insert(key, item)
+        else:
+            parent[key] = item
+    return state
+
+
+def gradio_measure(client, config, label, out):
+    fn = next(d for d in config['dependencies'] if d.get('api_name') == 'stage_perform')
+    session = uuid.uuid4().hex
+    started = time.perf_counter()
+    payload = {'fn_index': fn['id'], 'session_hash': session,
+               'data': [TEXT, 'Surprise', 'clone1', '', 'Narrator', '', '', 1, .9, 0, None, []]}
+    joined = client.post(SPACE + '/gradio_api/queue/join', json=payload)
+    joined.raise_for_status()
+    join_ms = (time.perf_counter() - started) * 1000
+    event_id = joined.json()['event_id']
+    state, first, prepared, messages, decoded_bytes = None, None, None, 0, 0
+    with client.stream('GET', SPACE + '/gradio_api/queue/data', params={'session_hash': session}) as response:
+        response.raise_for_status()
+        for line in response.iter_lines():
+            decoded_bytes += len(line)
+            if not line.startswith('data: '):
+                continue
+            event = json.loads(line[6:])
+            if event.get('event_id') != event_id:
+                continue
+            if event.get('success') is False:
+                raise RuntimeError(str(event.get('output')))
+            output = event.get('output', {}).get('data')
+            if output and event['msg'] in {'process_generating', 'process_completed'}:
+                messages += 1
+                state = apply_diff(state, output[0])
+                if state.get('type') == 'error':
+                    raise RuntimeError(state['message'])
+                if prepared is None:
+                    prepared = time.perf_counter() - started
+                if first is None and state.get('packets'):
+                    first = time.perf_counter() - started
+            if event['msg'] == 'process_completed':
+                break
+        wire_bytes = response.num_bytes_downloaded
+        encoding, version = response.headers.get('content-encoding', 'identity'), response.http_version
+    if not state or state.get('type') != 'done':
+        raise RuntimeError('The Gradio stage did not complete')
+    pcm = b''.join(base64.b64decode(p) for p in state['packets'])
+    with wave.open(str(out / (label + '.wav')), 'wb') as audio:
+        audio.setnchannels(1); audio.setsampwidth(2); audio.setframerate(state['sample_rate']); audio.writeframes(pcm)
+    return {'label': label, 'target': 'gradio', 'join_ms': round(join_ms),
+            'preparing_ms': round(prepared * 1000), 'first_audio_ms': round(first * 1000),
+            'wall_s': round(time.perf_counter() - started, 3), 'messages': messages,
+            'wire_bytes': wire_bytes, 'decoded_bytes': decoded_bytes,
+            'http_version': version, 'content_encoding': encoding,
+            **{key: state.get(key) for key in ['ttfa_ms', 'total_ms', 'timings_ms', 'audio_s', 'finish_reason']}}
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--out', type=Path, required=True)
+    parser.add_argument('--repeats', type=int, default=3)
+    parser.add_argument('--compression', choices=['gzip', 'identity'], default='gzip')
+    args = parser.parse_args()
+    args.out.mkdir(parents=True, exist_ok=True)
+    api = HfApi()
+    endpoint = api.get_inference_endpoint('drama-3-h200', namespace='HuggingFaceM4')
+    headers = {'Authorization': 'Bearer ' + get_token(), 'Accept-Encoding': args.compression}
+    rows = []
+    with httpx.Client(headers=headers, http2=True,
+                      limits=httpx.Limits(keepalive_expiry=120), timeout=180) as client:
+        config = client.get(SPACE + '/config').raise_for_status().json()
+        metadata = {'space_revision': api.space_info('andito/Drama-3').sha,
+                    'endpoint_args': endpoint.raw['model']['args'], 'settings': SETTINGS,
+                    'text': TEXT, 'repeats': args.repeats, 'accept_encoding': args.compression}
+        (args.out / 'metadata.json').write_text(json.dumps(metadata, indent=2) + '\n')
+        for repeat in range(args.repeats + 1):
+            order = ['endpoint', 'proxy', 'gradio'] if repeat % 2 == 0 else ['gradio', 'proxy', 'endpoint']
+            for target in order:
+                label = f'{target}-{repeat}'
+                if target == 'gradio':
+                    row = gradio_measure(client, config, label, args.out)
+                else:
+                    payload = {'input': TEXT, 'stream': True, 'response_format': 'pcm', 'stream_format': 'sse',
+                               'max_new_tokens': 500, **SETTINGS} if target == 'endpoint' else {
+                               'text': TEXT, 'voice': {'mode': 'random'}, **SETTINGS}
+                    url = endpoint.url + '/v1/audio/speech' if target == 'endpoint' else SPACE + '/api/generate'
+                    row = measure(client, url, payload, label, args.out)
+                    row['target'] = target
+                row.update(repeat=repeat, warmup=(repeat == 0))
+                rows.append(row)
+                print(json.dumps({k: v for k, v in row.items() if k != 'chunks'}), flush=True)
+                (args.out / 'results.json').write_text(json.dumps(rows, indent=2) + '\n')
+    summary = {}
+    for target in ['endpoint', 'proxy', 'gradio']:
+        samples = [r for r in rows if r['target'] == target and not r['warmup']]
+        summary[target] = {k: round(statistics.median(r[k] for r in samples), 3)
+                           for k in ['first_audio_ms', 'wall_s', 'wire_bytes', 'audio_s']}
+    (args.out / 'summary.json').write_text(json.dumps(summary, indent=2) + '\n')
+    print(json.dumps(summary, indent=2), flush=True)
+
+
+if __name__ == '__main__':
+    main()

endpoint/compare.py

@@ -0,0 +1,169 @@
+"""Repeat matched scripts on the H200 endpoint and the public Space; keep all audio."""
+import argparse
+import ast
+import base64
+import hashlib
+import json
+import math
+from pathlib import Path
+import re
+import time
+from types import SimpleNamespace
+
+import httpx
+import numpy as np
+import soundfile as sf
+from huggingface_hub import HfApi, get_token
+
+SPACE = 'https://andito-drama-3.hf.space'
+CASES = [
+    {'id': 'neutral-surprise', 'text': 'Good evening. The lights are on, the room is quiet, and we are ready to begin.', 'voice': {'mode': 'random'}},
+    {'id': 'expressive-surprise', 'text': "[sighs] Fine. If we miss it, it is on you. [laughs] Well, maybe that is fair.", 'voice': {'mode': 'random'}},
+    {'id': 'long-surprise', 'text': 'The rain has stopped, and the street is quiet. I left a note on the kitchen table. Please read it before you leave, then call me when you reach the station.', 'voice': {'mode': 'random'}},
+    {'id': 'clone1', 'text': 'Good evening. The lights are on, the room is quiet, and we are ready to begin.', 'voice': {'mode': 'clone', 'id': 'clone1'}},
+    {'id': 'design', 'text': '[softly] I thought you had forgotten. [sighs] I am glad you came back.', 'voice': {'mode': 'design', 'name': 'Narrator', 'description': 'Warm and quiet, a calm adult voice.'}},
+    {'id': 'whisper-clone3', 'text': "Promise me <whisper>you will not tell her</whisper>. [pause] Not a word. Not even a hint.", 'voice': {'mode': 'clone', 'id': 'clone3'}},
+]
+# top_k 4097 is no cap: every main-token candidate. top-k 30 makes takes go silent.
+SETTINGS = {'temperature': 1.0, 'top_p': .9, 'top_k': 4097}
+
+
+def helpers(source):
+    """Use the deployed Space's own text/budget functions, without importing its UI."""
+    tree = ast.parse(source)
+    names = {'spoken_text', 'token_budget'}
+    segments = [ast.get_source_segment(source, node) for node in tree.body
+                if isinstance(node, ast.FunctionDef) and node.name in names]
+    env = {'re': re, 'math': math, 'MAX_NEW_TOKENS': 1200, 'SAMPLE_RATE': 44100,
+           'DIRECTION_SPAN': re.compile(r'\[[^\[\]\n]*\]|</?[A-Za-z][^<>\n]{0,40}>')}
+    exec('from __future__ import annotations\n' + '\n\n'.join(segments), env)
+    return env['spoken_text'], env['token_budget']
+
+
+def audio_metrics(pcm, rate):
+    audio = np.frombuffer(pcm, dtype='<i2').astype(np.float32) / 32768
+    frame = round(.05 * rate)
+    quiet = [np.sqrt(np.mean(audio[i:i+frame] ** 2)) < .003 for i in range(0, len(audio), frame)]
+    speech = [i for i, q in enumerate(quiet) if not q]
+    leading = speech[0] * .05 if speech else len(quiet) * .05
+    trailing = (len(quiet) - speech[-1] - 1) * .05 if speech else len(quiet) * .05
+    runs, run = [], 0
+    internal = quiet[speech[0]:speech[-1]+1] if speech else []
+    for q in internal + [False]:
+        if q:
+            run += 1
+        elif run:
+            runs.append(run * .05)
+            run = 0
+    return {'audio_s': round(len(audio) / rate, 3), 'leading_quiet_s': round(leading, 2),
+            'trailing_quiet_s': round(trailing, 2), 'longest_internal_quiet_s': round(max(runs, default=0), 2),
+            'internal_gaps_over_2s': sum(r > 2 for r in runs),
+            'rms': round(float(np.sqrt(np.mean(audio ** 2))), 5) if len(audio) else 0}
+
+
+def generate(client, source, endpoint_url, case, native, out):
+    started = time.perf_counter()
+    audio = bytearray()
+    result = {'source': source, 'case': case['id'], 'repeat': case['repeat'], 'expected': case['expected'],
+              'max_new_tokens': native['max_new_tokens'], 'retry_count': 0}
+    url = endpoint_url + '/v1/audio/speech' if source == 'endpoint' else SPACE + '/api/generate'
+    payload = native if source == 'endpoint' else {'text': case['text'], 'voice': case['voice'], **SETTINGS}
+    rate, first = 44100, None
+    with client.stream('POST', url, json=payload) as response:
+        response.raise_for_status()
+        rate = int(response.headers.get('x-sample-rate', rate))
+        for raw in response.iter_lines():
+            if not raw.startswith('data: '):
+                continue
+            event = json.loads(raw[6:])
+            kind = event.get('type')
+            if kind == 'start':
+                rate = event['sample_rate']
+            if kind == 'retrying':
+                result['retry_count'] += 1
+            if kind in {'audio', 'speech.audio.delta'}:
+                if first is None:
+                    first = time.perf_counter()
+                audio.extend(base64.b64decode(event['pcm'] if source == 'space' else event['audio']))
+            if kind in {'done', 'speech.audio.done'}:
+                result['finish_reason'] = event.get('finish_reason')
+                if source == 'space':
+                    result.update({k: event[k] for k in ['warning', 'retry_count', 'trimmed_s', 'ttfa_ms', 'queue_ms'] if k in event})
+                break
+            if kind == 'error':
+                result['error'] = event.get('message', str(event))
+                break
+    result['wall_s'] = round(time.perf_counter() - started, 3)
+    result['first_audio_client_ms'] = round((first - started) * 1000) if first else None
+    result.update(audio_metrics(bytes(audio), rate))
+    if audio:
+        filename = f"{case['id']}-{case['repeat']}-{source}.wav"
+        sf.write(out / filename, np.frombuffer(audio, dtype='<i2'), rate, subtype='PCM_16')
+        result['audio_file'] = filename
+    return result
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--out', type=Path, required=True)
+    parser.add_argument('--repeats', type=int, default=4)
+    args = parser.parse_args()
+    out = args.out
+    out.mkdir(parents=True, exist_ok=True)
+    api = HfApi()
+    endpoint = api.get_inference_endpoint('drama-3-h200', namespace='HuggingFaceM4')
+    if endpoint.status != 'running' or not endpoint.url:
+        raise RuntimeError('Endpoint is not running')
+    space = api.space_info('andito/Drama-3')
+    model = api.model_info('fishaudio/drama-3')
+    if endpoint.raw['model']['revision'] != model.sha:
+        raise RuntimeError('Model revisions differ; resolve this before comparing')
+    results = []
+    with httpx.Client(headers={'Authorization': 'Bearer ' + get_token()}, timeout=120, follow_redirects=True) as client:
+        status = client.get(SPACE + '/api/status').raise_for_status().json()
+        assert status['engine']['state'] == 'ready' and status['load']['active'] == 0
+        assert status['transcriber'] == 'ready'
+        code = client.get(f'https://huggingface.co/spaces/andito/Drama-3/resolve/{space.sha}/server.py').raise_for_status().text
+        spoken, budget = helpers(code)
+        opening_buffer = float(re.search(r'^START_BUFFER_S\s*=\s*([\d.]+)', code, re.M).group(1))
+        voices = {v['id']: v for v in status['voices']}
+        refs = {}
+        for name in ['clone1', 'clone3']:
+            content = client.get(SPACE + '/api/voices/' + name + '.wav').raise_for_status().content
+            refs[name] = content
+        metadata = {'space_revision': space.sha, 'endpoint_model_revision': endpoint.raw['model']['revision'],
+                    'space_gpu': status['gpu'], 'endpoint_compute': endpoint.raw['compute'],
+                    'endpoint_image': endpoint.raw['model']['image'], 'settings': SETTINGS,
+                    'reference_sha256': {k: hashlib.sha256(v).hexdigest() for k, v in refs.items()},
+                    'repeats': args.repeats,
+                    'opening_buffer_s': opening_buffer,
+                    'note': f'Space includes trimming, a {opening_buffer:g}-second opening buffer, and up to one retry; native endpoint audio is unprocessed. No seed is exposed by the Space API. Balanced sequential order; each target has the same model revision, script, reference audio, text budget, and sampling settings.'}
+        (out / 'metadata.json').write_text(json.dumps(metadata, indent=2) + '\n')
+        for repeat in range(1, args.repeats + 1):
+            for index, template in enumerate(CASES):
+                case = {**template, 'repeat': repeat, 'expected': spoken(template['text'])}
+                voice = case['voice']
+                char = SimpleNamespace(description=voice['description'], scene=voice.get('scene')) if voice['mode'] == 'design' else None
+                tokens = budget(SimpleNamespace(text=case['text'], note=None), SimpleNamespace(character=char))
+                input_text = case['text']
+                references = None
+                if voice['mode'] == 'design':
+                    input_text = '@@drama3-space@@' + json.dumps({'text': case['text'], 'voice_references': 0, 'history': [],
+                               'character': {'name': voice['name'], 'description': voice['description'], 'scene': voice.get('scene')}}, ensure_ascii=False)
+                if voice['mode'] == 'clone':
+                    references = [{'data': base64.b64encode(refs[voice['id']]).decode(), 'media_type': 'audio/wav', 'text': voices[voice['id']]['transcript']}]
+                native = {'input': input_text, 'references': references, 'stream': True, 'response_format': 'pcm',
+                          'stream_format': 'sse', 'max_new_tokens': tokens, **SETTINGS}
+                order = ['space', 'endpoint'] if (repeat + index) % 2 else ['endpoint', 'space']
+                for source in order:
+                    result = generate(client, source, endpoint.url.rstrip('/'), case, native, out)
+                    results.append(result)
+                    print(json.dumps({k:result.get(k) for k in ['case','repeat','source','finish_reason','retry_count','audio_s','longest_internal_quiet_s','trailing_quiet_s','error']}), flush=True)
+                    (out / 'generations.json').write_text(json.dumps(results, indent=2) + '\n')
+        print('All matched generations saved', flush=True)
+        status = client.get(SPACE + '/api/status').raise_for_status().json()
+        print('Final Space load:', status['load'], flush=True)
+
+
+if __name__ == '__main__':
+    main()

endpoint/deploy.py

@@ -0,0 +1,51 @@
+"""Create the single-GPU Drama 3 inference server; reuse matching deployments."""
+import argparse
+import json
+from pathlib import Path
+from huggingface_hub import HfApi
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--update-serving', action='store_true',
+                        help='Apply the pinned container arguments to an existing matching endpoint.')
+    args = parser.parse_args()
+    config = json.loads(Path(__file__).with_name("deployment.json").read_text())
+    api = HfApi()
+    existing = next((e for e in api.list_inference_endpoints(namespace=config["namespace"])
+                     if e.name == config["name"]), None)
+    if existing is not None:
+        compute = existing.raw["compute"]
+        if (existing.repository != config["repository"]
+                or existing.raw["provider"] != {"vendor": config["vendor"], "region": config["region"]}
+                or compute["instanceType"] != config["instance_type"]
+                or compute["instanceSize"] != config["instance_size"]):
+            raise RuntimeError("The endpoint name belongs to a different configuration; refusing to overwrite it.")
+        endpoint = existing
+        if args.update_serving:
+            model = existing.raw['model']
+            image = model.get('image', {}).get('custom', {})
+            pinned_image = config['custom_image']
+            if (model.get('revision') != config['revision']
+                    or image.get('url') != pinned_image['url']
+                    or image.get('port') != pinned_image['port']
+                    or image.get('healthRoute') != pinned_image['health_route']
+                    or model.get('command') != config['container_command']):
+                raise RuntimeError('Model revision, image, or command differ; review before updating serving flags.')
+            if model.get('args') != config['container_args']:
+                print('Updating serving arguments on the existing endpoint; hardware and scaling stay the same.', flush=True)
+                endpoint = existing.update(container_args=config['container_args'])
+    else:
+        hardware = api.list_inference_endpoints_hardware(namespace=config["namespace"])
+        target = next((h for h in hardware if (h.vendor, h.region, h.instance_type, h.instance_size)
+                       == (config["vendor"], config["region"], config["instance_type"], config["instance_size"])), None)
+        if target is None or target.status != "available" or target.max_accelerators - target.used_accelerators < 1:
+            raise RuntimeError("A single GCP H200 is unavailable or exceeds this organization's quota.")
+        print(f"Creating {config['namespace']}/{config['name']}: ${target.price_per_hour}/hour while running", flush=True)
+        endpoint = api.create_inference_endpoint(**config)
+    print(json.dumps({"name": endpoint.name, "namespace": endpoint.namespace,
+                      "status": endpoint.status, "url": endpoint.url}), flush=True)
+
+
+if __name__ == "__main__":
+    main()

endpoint/deployment.json

@@ -0,0 +1,50 @@
+{
+  "name": "drama-3-h200",
+  "namespace": "HuggingFaceM4",
+  "repository": "fishaudio/drama-3",
+  "revision": "e0a9218aa48888270f06ef1ed9ea67622c0e17d2",
+  "framework": "custom",
+  "task": "custom",
+  "accelerator": "gpu",
+  "instance_size": "x1",
+  "instance_type": "nvidia-h200",
+  "region": "us-south1",
+  "vendor": "gcp",
+  "type": "authenticated",
+  "min_replica": 0,
+  "max_replica": 1,
+  "scale_to_zero_timeout": 60,
+  "custom_image": {
+    "url": "registry.hf.space/andito-drama-3@sha256:e4d8bfb5af1f32029a36d7e9d1bafd83d9416cdc15915f71dce14b1853c42d0f",
+    "port": 8000,
+    "health_route": "/health"
+  },
+  "container_command": [
+    "python3",
+    "-m",
+    "sglang_omni.cli"
+  ],
+  "container_args": [
+    "serve",
+    "--model-path",
+    "/repository",
+    "--host",
+    "0.0.0.0",
+    "--port",
+    "8000",
+    "--tts_engine.factory.max_running_requests",
+    "16",
+    "--tts_engine.factory.max_new_tokens",
+    "1200",
+    "--tts_engine.factory.enable_async_decode",
+    "true",
+    "--tts_engine.factory.async_decode_min_batch_size",
+    "1"
+  ],
+  "tags": [
+    "drama-3",
+    "tts",
+    "gcp",
+    "h200"
+  ]
+}

endpoint/verify.py

@@ -0,0 +1,71 @@
+"""Verify streaming speech on the deployed server; save audio and metrics locally."""
+import base64
+import io
+import json
+import time
+from pathlib import Path
+import httpx
+import numpy as np
+import soundfile as sf
+from huggingface_hub import HfApi, get_token
+
+
+def main():
+    root = Path(__file__).parent
+    config = json.loads((root / 'deployment.json').read_text())
+    endpoint = HfApi().get_inference_endpoint(config['name'], namespace=config['namespace'])
+    if endpoint.status != 'running' or not endpoint.url:
+        raise RuntimeError('The endpoint must be running before verifying speech')
+    result = []
+    cases = [
+        ('plain', 'Good evening. The lights are on, the room is quiet, and we are ready to begin.', None),
+        ('expressive', '[softly] I thought you had forgotten. [sighs] I am glad you came back.', None),
+        ('clone', 'A new stage, a familiar voice. This is Drama Three on an H two hundred.', 'clone3.wav'),
+    ]
+    with httpx.Client(headers={'Authorization': 'Bearer ' + get_token()}, timeout=120) as client:
+        health = client.get(endpoint.url + '/health')
+        health.raise_for_status()
+        for name, text, reference in cases:
+            payload = {'input': text, 'stream': True, 'response_format': 'pcm', 'stream_format': 'sse',
+                       'max_new_tokens': 480, 'temperature': 1.0, 'top_p': .9, 'top_k': 4097}
+            if reference:
+                voice_root = root.parent / 'voices'
+                entry = next(v for v in json.loads((voice_root / 'voices.json').read_text()) if v['file'] == reference)
+                payload['references'] = [{'data': base64.b64encode((voice_root / reference).read_bytes()).decode(),
+                                          'media_type': 'audio/wav', 'text': entry['transcript']}]
+            audio = bytearray()
+            first = None
+            finish = None
+            rate = 44100
+            started = time.perf_counter()
+            with client.stream('POST', endpoint.url + '/v1/audio/speech', json=payload) as response:
+                response.raise_for_status()
+                rate = int(response.headers.get('x-sample-rate', rate))
+                for raw in response.iter_lines():
+                    if not raw.startswith('data: '):
+                        continue
+                    event = json.loads(raw[6:])
+                    if event.get('type') == 'error':
+                        raise RuntimeError(event)
+                    if event.get('type') == 'speech.audio.delta':
+                        if first is None:
+                            first = time.perf_counter()
+                        audio.extend(base64.b64decode(event['audio']))
+                    if event.get('type') == 'speech.audio.done':
+                        finish = event.get('finish_reason')
+                        break
+            samples = np.frombuffer(audio, dtype='<i2').astype(np.float32) / 32768
+            if samples.size == 0 or float(np.sqrt(np.mean(samples ** 2))) < .003:
+                raise RuntimeError(f'{name} produced no audible speech')
+            sf.write(root / (name + '.wav'), samples, rate, subtype='PCM_16')
+            elapsed = time.perf_counter() - started
+            metrics = {'case': name, 'finish_reason': finish, 'audio_s': round(samples.size / rate, 2),
+                       'first_audio_ms': round((first - started) * 1000), 'wall_time_s': round(elapsed, 2)}
+            result.append(metrics)
+            print(json.dumps(metrics), flush=True)
+            (root / 'verification.json').write_text(json.dumps(result, indent=2) + '\n')
+    print('Speech streaming and voice cloning verified', flush=True)
+
+
+if __name__ == '__main__':
+    main()

endpoint/watch.py

@@ -0,0 +1,38 @@
+"""Watch this deployment without exposing credentials or endpoint secrets."""
+import json
+import time
+from pathlib import Path
+import httpx
+from huggingface_hub import HfApi, get_token
+
+
+def main():
+    root = Path(__file__).parent
+    config = json.loads((root / 'deployment.json').read_text())
+    api = HfApi()
+    last = None
+    deadline = time.monotonic() + 2400
+    with httpx.Client(headers={'Authorization': 'Bearer ' + get_token()}, timeout=30) as client:
+        while time.monotonic() < deadline:
+            endpoint = api.get_inference_endpoint(config['name'], namespace=config['namespace'])
+            state = str(endpoint.status)
+            summary = {'name': endpoint.name, 'namespace': endpoint.namespace, 'status': state,
+                       'url': endpoint.url, 'provider': endpoint.raw.get('provider'),
+                       'compute': endpoint.raw.get('compute'), 'deployment_status': endpoint.raw.get('status')}
+            (root / 'status.json').write_text(json.dumps(summary, indent=2) + '\n')
+            if state != last:
+                print(json.dumps(summary), flush=True)
+                last = state
+            if state in {'failed', 'updateFailed'}:
+                raise RuntimeError('Endpoint deployment failed; inspect the deployment status and logs')
+            if state == 'running' and endpoint.url:
+                response = client.get(endpoint.url.rstrip('/') + '/health')
+                if response.status_code == 200:
+                    print('Speech server health check passed', flush=True)
+                    return
+            time.sleep(30)
+    raise RuntimeError('Timed out waiting for the speech server')
+
+
+if __name__ == '__main__':
+    main()

entitlements.py

@@ -0,0 +1,134 @@
+"""Resolve demo quota tiers from Hugging Face OAuth identity."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import hashlib
+import logging
+import os
+import threading
+from typing import Any, Callable
+
+
+logger = logging.getLogger("drama3_space.entitlements")
+
+_DEFAULT_UNLIMITED_ORGS = {"huggingfacem4"}
+_profile_cache: dict[str, dict[str, Any]] = {}
+_profile_cache_lock = threading.Lock()
+
+
+@dataclass(frozen=True)
+class Entitlement:
+    tier: str
+    is_pro: bool
+    is_team_member: bool
+
+    @property
+    def unlimited(self) -> bool:
+        return self.is_pro or self.is_team_member
+
+
+def _field(obj: Any, name: str, default: Any = None) -> Any:
+    if obj is None:
+        return default
+    if isinstance(obj, dict):
+        return obj.get(name, default)
+    return getattr(obj, name, default)
+
+
+def _unlimited_orgs() -> set[str]:
+    """Return case-insensitive org names whose members have unlimited usage."""
+    raw = " ".join(
+        (
+            os.environ.get("UNLIMITED_ORGS", ""),
+            os.environ.get("DEMO_UNLIMITED_ORGS", ""),
+        )
+    )
+    configured = {
+        value.strip().lower()
+        for value in raw.replace(",", " ").split()
+        if value.strip()
+    }
+    return _DEFAULT_UNLIMITED_ORGS | configured
+
+
+def _oauth_token(oauth_info: Any) -> str | None:
+    token = _field(oauth_info, "access_token")
+    if token is None:
+        return None
+    return str(token).strip() or None
+
+
+def _profile_cache_key(token: str) -> str:
+    return hashlib.sha256(token.encode("utf-8")).hexdigest()
+
+
+def _whoami_via_token(token: str) -> dict[str, Any]:
+    """Fetch and cache the authenticated Hub ``whoami-v2`` profile."""
+    cache_key = _profile_cache_key(token)
+    with _profile_cache_lock:
+        cached = _profile_cache.get(cache_key)
+    if cached is not None:
+        return cached
+
+    try:
+        import httpx
+
+        response = httpx.get(
+            "https://huggingface.co/api/whoami-v2",
+            headers={"Authorization": f"Bearer {token}"},
+            timeout=5.0,
+        )
+        response.raise_for_status()
+        profile = response.json()
+        if not isinstance(profile, dict):
+            raise ValueError("whoami-v2 returned a non-object response")
+    except Exception as exc:  # pragma: no cover - network/permission dependent
+        logger.info("whoami-v2 profile lookup failed: %r", exc)
+        return {}
+
+    with _profile_cache_lock:
+        _profile_cache[cache_key] = profile
+    return profile
+
+
+def _org_names(user: Any, profile: dict[str, Any]) -> set[str]:
+    names: set[str] = set()
+    org_sources = (_field(user, "orgs", []) or [], profile.get("orgs", []) or [])
+    for orgs in org_sources:
+        for org in orgs:
+            for key in ("preferred_username", "name", "fullname", "sub"):
+                value = _field(org, key)
+                if value:
+                    names.add(str(value).lower())
+    return names
+
+
+_PAID_ORG_PLANS = {"team", "enterprise", "enterprise_plus"}
+
+
+def _in_paid_org(user: Any, profile: dict[str, Any]) -> bool:
+    """Member of an org on a Team or Enterprise plan ("PRO and up")."""
+    for orgs in (_field(user, "orgs", []) or [], profile.get("orgs", []) or []):
+        for org in orgs:
+            if _field(org, "isEnterprise") or str(_field(org, "plan", "")).lower() in _PAID_ORG_PLANS:
+                return True
+    return False
+
+
+def resolve_entitlement(
+    oauth_info: Any,
+    *,
+    profile_lookup: Callable[[str], dict[str, Any]] | None = None,
+) -> Entitlement:
+    """Resolve ``free``, ``pro``, or ``team`` from a Space OAuth session."""
+    user = _field(oauth_info, "user_info")
+    token = _oauth_token(oauth_info)
+    profile: dict[str, Any] = {}
+    if token:
+        profile = (profile_lookup or _whoami_via_token)(token)
+
+    is_pro = bool(_field(user, "is_pro", False) or _field(profile, "isPro", False))
+    is_team_member = bool(_unlimited_orgs() & _org_names(user, profile)) or _in_paid_org(user, profile)
+    tier = "pro" if is_pro else "team" if is_team_member else "free"
+    return Entitlement(tier=tier, is_pro=is_pro, is_team_member=is_team_member)

fish_speech/LICENSE

@@ -1,114 +0,0 @@
-# FISH AUDIO COMMUNITY LICENSE AGREEMENT
-
-**Version 1.0 — Last Updated: October 13, 2026**
-
-## I. INTRODUCTION
-
-This Agreement applies to any individual person or entity ("You", "Your" or "Licensee") that uses or distributes any portion or element of the Fish Audio Materials or Derivative Works thereof for any Research, Non-Commercial, or Commercial purpose. Capitalized terms not otherwise defined herein are defined in Section V below.
-
-This Agreement is intended to allow, free of charge, research and non-commercial uses of the Fish Audio Materials and, as set out in Section III, limited uses for a Commercial Purpose by Licensees whose Annual Revenue is less than the Revenue Threshold. Any other use of the Fish Audio Materials for a Commercial Purpose, and any Hosted Service, requires a separate written license agreement from Fish Audio.
-
-By clicking "I Accept" or by using, distributing, or accessing any portion or element of the Fish Audio Materials or Derivative Works, You agree that You have read, understood and are bound by the terms of this Agreement. If You are acting on behalf of a company, organization or other entity, then "You" includes you and that entity, and You agree that You: (i) are an authorized representative of such entity with the authority to bind such entity to this Agreement, and (ii) You agree to the terms of this Agreement on that entity's behalf.
-
-## II. RESEARCH & NON-COMMERCIAL USE LICENSE
-
-Subject to the terms of this Agreement, Fish Audio grants You a non-exclusive, worldwide, non-transferable, non-sublicensable, revocable and royalty-free limited license under Fish Audio's intellectual property or other rights owned by Fish Audio embodied in the Fish Audio Materials to use, reproduce, distribute, and create Derivative Works of, and make modifications to, the Fish Audio Materials for any Research or Non-Commercial Purpose.
-
-"Research Purpose" means academic or scientific advancement, and in each case, is not primarily intended for commercial advantage or monetary compensation to You or others.
-
-"Non-Commercial Purpose" means any purpose other than a Research Purpose that is not primarily intended for commercial advantage or monetary compensation to You or others, such as personal use (i.e., hobbyist) or evaluation and testing.
-
-## III. COMMERCIAL USE
-
-**(a) Community Commercial License.** Subject to the terms of this Agreement (including Sections III(b) through III(d)), Fish Audio grants You a non-exclusive, worldwide, non-transferable, non-sublicensable, revocable and royalty-free limited license under Fish Audio's intellectual property or other rights owned by Fish Audio embodied in the Fish Audio Materials to use, reproduce, distribute, and create Derivative Works of, and make modifications to, the Fish Audio Materials for a Commercial Purpose, but only while Your Annual Revenue is less than USD 500,000 (the "Revenue Threshold").
-
-"Annual Revenue" means, as of any date, the total gross revenue of You and Your Affiliates, in aggregate, from all sources (whether or not derived from the Fish Audio Materials or Derivative Works) for the twelve (12) months ending on that date, with amounts in other currencies converted into U.S. dollars at a published market exchange rate in effect on that date. If You deploy, install, host, or operate the Fish Audio Materials or Derivative Works for a particular third party, or deliver a copy of them to a third party, for use in that third party's products, services, or internal operations (for example, as a contractor, integrator, consultant, or on-premises vendor), the license in this Section III(a) extends to that work only if that third party's Annual Revenue, calculated as if that third party were You, is also less than the Revenue Threshold. This does not apply to (1) delivering to a third party only output (such as audio files) that You generate Yourself, for example as an agency, studio, or freelancer producing voiceovers, dubbing, or audiobooks for a client, or (2) providing Your own generally available product or service to customers or end users who do not receive a copy of the Fish Audio Materials or Derivative Works; such customers and end users do not need to meet the Revenue Threshold for that use.
-
-**(b) Reaching the Revenue Threshold.** If, while You are using the Fish Audio Materials or Derivative Works under Section III(a), Your Annual Revenue reaches or exceeds the Revenue Threshold (the date it first does so, the "Threshold Date"), You must notify Fish Audio in writing at business@fish.audio within thirty (30) days after the Threshold Date. Notwithstanding Section III(a), the license in Section III(a) continues to apply to the uses for a Commercial Purpose that You were making on the Threshold Date for ninety (90) days after the Threshold Date (the "Transition Period"). After the Transition Period, any use for a Commercial Purpose requires a separate written license agreement from Fish Audio; however, if Your Annual Revenue later falls below the Revenue Threshold, Section III(a) applies to You again from that date. Reaching the Revenue Threshold is not a breach of this Agreement and does not affect Your license under Section II or Your ownership of outputs under Section IV(c)(iii).
-
-**(c) Hosted Services.** Notwithstanding Sections II and III(a) and the definition of Commercial Purpose, and regardless of Your purpose or Annual Revenue, no license is granted under this Agreement to offer a Hosted Service, which requires a separate written license agreement from Fish Audio. "Hosted Service" means making the Fish Audio Materials or Derivative Works available to third parties so that they can generate speech or other audio from text, audio, or voice samples that they supply, whether through an application programming interface, software development kit, model-inference or model-hosting service, or a website, application, or feature whose primary purpose is text-to-speech, voiceover, voice generation, or voice cloning. None of the following is a Hosted Service: (1) generating speech for content that You create, select, or control, or that Your own product generates (for example, narration, dubbing, game or app characters, or an assistant's spoken replies), and delivering that speech in or with Your products or services; (2) a production service in which You operate the Fish Audio Materials or Derivative Works, direct each generation, and deliver only finished output to Your client (for example, as an agency, studio, or freelancer producing voiceovers, dubbing, or audiobooks), without giving the client self-serve or automated access; or (3) a free, non-commercial research or demonstration service (for example, a Hugging Face Space, a research demo, or a shared server for a research group) that is not monetized and is not offered or promoted as a production service or API.
-
-**(d) Other Commercial Use.** **Any use of the Fish Audio Materials or Derivative Works for a Commercial Purpose that is not licensed under Section III(a) (including as continued under Section III(b)) requires a separate written license agreement from Fish Audio.**
-
-"Commercial Purpose" means any purpose other than a Research Purpose or Non-Commercial Purpose that is primarily intended for or directed toward commercial advantage or monetary compensation to You or others, including but not limited to: (i) creating, modifying, or distributing Your product or service, including via a hosted service or application programming interface, (ii) Your business's or organization's internal operations, and (iii) any use in connection with a product or service for which You charge a fee or generate revenue, whether directly or indirectly.
-
-To obtain a commercial license, please contact Fish Audio at:
-
-- **Website:** [https://fish.audio](https://fish.audio)
-- **Email:** business@fish.audio
-
-## IV. GENERAL TERMS
-
-All licenses granted under this Agreement, including Your Research, Non-Commercial, and Commercial licenses, are subject to the following terms.
-
-### a. Distribution & Attribution
-
-If You distribute a copy of the Fish Audio Materials or a Derivative Work to a third party, You shall (i) provide a copy of this Agreement to that third party, and (ii) retain the following attribution notice within a "Notice" text file distributed as a part of such copies: "These Fish Audio Materials are licensed under the Fish Audio Community License, Version 1.0. Copyright © 39 AI, INC. All Rights Reserved." If You distribute or make available the Fish Audio Materials or a Derivative Work, or a product or service that uses any portion of them, You shall prominently display "Built with Fish Audio" on a related website, user interface, blogpost, about page, or product documentation.
-
-If You create a Derivative Work, You may add your own attribution notice(s) to the "Notice" text file included with that Derivative Work, provided that You clearly indicate which attributions apply to the Fish Audio Materials and state in the "Notice" text file that You changed the Fish Audio Materials and how it was modified.
-
-### b. Use Restrictions
-
-Your use of the Fish Audio Materials and Derivative Works, including any output or results of the Fish Audio Materials or Derivative Works, must comply with applicable laws and regulations (including Trade Control Laws and equivalent regulations).
-
-You will not use the Fish Audio Materials or Derivative Works, or any output or results of the Fish Audio Materials or Derivative Works, to: (i) generate speech that imitates the voice of an identifiable natural person without that person's prior consent (or, for a minor, the consent of a parent or legal guardian, and for a deceased person, the consent of the holder of the relevant rights) or another legal right to do so; (ii) impersonate any person or organization in order to deceive, defraud, or harass, or to place calls in violation of applicable telemarketing or robocall laws; (iii) generate sexual content involving a minor or a voice presented as a minor's, or using the voice of any real person without that person's consent; or (iv) intentionally remove, disable, circumvent, or obscure any watermark or provenance signal applied by the Fish Audio Materials (ordinary editing, mixing, compression, or format conversion of output that is not intended to remove such a signal is not a breach of this clause (iv)). You will not distribute any tool, Derivative Work, or service designed to remove, disable, circumvent, or obscure any such watermark or provenance signal. If You make voice cloning or other speech generation using the Fish Audio Materials or Derivative Works available to third parties, You must bind them, by enforceable terms, to restrictions at least as protective as those in this Section IV(b).
-
-Furthermore, You will not use the Fish Audio Materials or Derivative Works, or any output or results of the Fish Audio Materials or Derivative Works, to create, train, distill, or improve any foundational generative AI model or any speech, audio, or voice generation model (excluding the Models or Derivative Works). This does not restrict training models solely to detect synthetic audio, verify watermarks, or prevent spoofing.
-
-### c. Intellectual Property
-
-**(i) Trademark License.** No trademark licenses are granted under this Agreement. In connection with the Fish Audio Materials or Derivative Works, You may not use any name or mark owned by or associated with Fish Audio or any of its Affiliates, except (A) as required under Section IV(a), and (B) to make truthful, factual references to Fish Audio and the Models, including naming a Derivative Work model in a way that identifies its origin (for example, "Drama-3-GGUF" or "based on Fish Audio Drama 3"), provided that You do not use any Fish Audio logo, do not use a Fish Audio name or mark as the name of Your company or of any product or service other than such a Derivative Work model, and do not state or imply that Fish Audio sponsors, endorses, or officially released You, Your products, or Your Derivative Works.
-
-**(ii) Ownership of Derivative Works.** As between You and Fish Audio, You are the owner of Derivative Works You create, subject to Fish Audio's ownership of the Fish Audio Materials and any Derivative Works made by or for Fish Audio.
-
-**(iii) Ownership of Outputs.** As between You and Fish Audio, You own any outputs generated from the Models or Derivative Works to the extent permitted by applicable law.
-
-**(iv) Disputes.** If You or Your Affiliate(s) institute litigation or other proceedings against Fish Audio (including a cross-claim or counterclaim in a lawsuit) alleging that the Fish Audio Materials, Derivative Works or associated outputs or results, or any portion of any of the foregoing, constitutes infringement of intellectual property or other rights owned or licensable by You, then any licenses granted to You under this Agreement shall terminate as of the date such litigation or claim is filed or instituted. You will indemnify and hold harmless Fish Audio from and against any claim by any third party arising out of or related to Your use or distribution of the Fish Audio Materials or Derivative Works in violation of this Agreement.
-
-**(v) Feedback.** From time to time, You may provide Fish Audio with verbal and/or written suggestions, comments or other feedback related to Fish Audio's existing or prospective technology, products or services (collectively, "Feedback"). You are not obligated to provide Fish Audio with Feedback, but to the extent that You do, You hereby grant Fish Audio a perpetual, irrevocable, royalty-free, fully-paid, sub-licensable, transferable, non-exclusive, worldwide right and license to exploit the Feedback in any manner without restriction. Your Feedback is provided "AS IS" and You make no warranties whatsoever about any Feedback.
-
-### d. Disclaimer of Warranty
-
-UNLESS REQUIRED BY APPLICABLE LAW, THE FISH AUDIO MATERIALS AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OR LAWFULNESS OF USING OR REDISTRIBUTING THE FISH AUDIO MATERIALS, DERIVATIVE WORKS OR ANY OUTPUT OR RESULTS AND ASSUME ANY RISKS ASSOCIATED WITH YOUR USE OF THE FISH AUDIO MATERIALS, DERIVATIVE WORKS AND ANY OUTPUT AND RESULTS.
-
-### e. Limitation of Liability
-
-IN NO EVENT WILL FISH AUDIO OR ITS AFFILIATES BE LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, ARISING OUT OF THIS AGREEMENT, FOR ANY LOST PROFITS OR ANY DIRECT, INDIRECT, SPECIAL, CONSEQUENTIAL, INCIDENTAL, EXEMPLARY OR PUNITIVE DAMAGES, EVEN IF FISH AUDIO OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF ANY OF THE FOREGOING.
-
-### f. Term and Termination
-
-The term of this Agreement will commence upon Your acceptance of this Agreement or access to the Fish Audio Materials and will continue in full force and effect until terminated in accordance with the terms and conditions herein. Fish Audio may terminate this Agreement by written notice to You if You breach any term or condition of this Agreement and, if the breach is capable of cure, You do not cure it within thirty (30) days after Fish Audio's written notice describing the breach. If You breach Section IV(b), this Agreement and all licenses granted to You under it terminate automatically, without notice or cure period. Upon termination of this Agreement, You shall delete and cease use of any Fish Audio Materials or Derivative Works. Sections IV(b) (as it applies to any output or results of the Fish Audio Materials or Derivative Works that You retain), IV(c)(iii), IV(c)(iv), IV(c)(v), IV(d), IV(e), IV(g) and V shall survive the termination of this Agreement.
-
-### g. Governing Law
-
-This Agreement will be governed by and construed in accordance with the laws of the United States and the State of California without regard to choice of law principles, and the UN Convention on Contracts for International Sale of Goods does not apply to this Agreement.
-
-### h. Versions
-
-Fish Audio may publish revised versions of the Fish Audio Community License, each identified by a version number. Fish Audio Materials are licensed under the version with which Fish Audio released them, as stated in the license file accompanying them or on Fish Audio's official model page for them. A revised version applies only to Fish Audio Materials that Fish Audio releases or designates under it, and does not reduce the rights You hold under the version You received. You may elect to use Fish Audio Materials under a later version that Fish Audio designates as applicable to them.
-
-### i. Separate Agreements
-
-If You and Fish Audio have signed a separate written agreement covering the Fish Audio Materials, that agreement governs to the extent it conflicts with this Agreement.
-
-## V. DEFINITIONS
-
-**"Affiliate(s)"** means any entity that directly or indirectly controls, is controlled by, or is under common control with the subject entity; for purposes of this definition, "control" means direct or indirect ownership or control of more than 50% of the voting interests of the subject entity.
-
-**"Agreement"** means this Fish Audio Community License Agreement, Version 1.0 (also referred to as the "Fish Audio Community License, Version 1.0" or "FACL v1.0").
-
-**"Annual Revenue"** and **"Revenue Threshold"** have the meanings given in Section III(a), and **"Hosted Service"** has the meaning given in Section III(c).
-
-**"Derivative Work(s)"** means (a) any derivative work of the Fish Audio Materials as recognized by U.S. copyright laws and (b) any modification of a Model, including any model initialized from all or part of a Model's weights (such as fine-tuned, quantized, merged, and low-rank adaptation ("LoRA") models), whether or not it was trained using output of a Model or Derivative Work. Derivative Works do not include the output of any Model or Derivative Work, or any other model trained or distilled using such output.
-
-**"Documentation"** means any specifications, manuals, model cards, documentation, and other written information that Fish Audio makes available under this Agreement related to the Software or Models.
-
-**"Fish Audio"** or **"we"** means 39 AI, INC. and its Affiliates.
-
-**"Model(s)"** means the machine-learning models, including trained model weights and other elements of the foregoing, that Fish Audio makes available under this Agreement, including Fish Audio Drama 3 and, solely for use with Fish Audio Drama 3 and its Derivative Works, the audio codec that Fish Audio provides for use with Fish Audio Drama 3 (including the copy of that codec distributed with Fish Audio S2-Pro).
-
-**"Software"** means Fish Audio's proprietary software made available under this Agreement now or in the future, including the Fish Speech source code (https://github.com/fishaudio/fish-speech).
-
-**"Fish Audio Materials"** means, collectively, the Models, Software and Documentation (and any portion or combination thereof) made available under this Agreement. Model weights that Fish Audio makes available under a different license (for example, the Fish Audio S2-Pro weights under the Fish Audio Research License) are not Fish Audio Materials, even if the Software can run them, except as stated in the definition of Model(s) for the audio codec.
-
-**"Trade Control Laws"** means any applicable U.S. and non-U.S. export control and trade sanctions laws and regulations.

fish_speech/__init__.py

@@ -1 +0,0 @@
-

fish_speech/codec.py

@@ -1,113 +0,0 @@
-"""Builds the Drama 3 / S2-Pro codec (configs/modded_dac_vq.yaml) without Hydra,
-and mirrors load_codec_model / encode_audio / decode_to_audio from
-fish_speech/models/text2semantic/inference.py."""
-
-from functools import partial
-
-import torch
-from torch.nn.utils.parametrize import is_parametrized, remove_parametrizations
-
-from fish_speech.models.dac.modded_dac import DAC, ModelArgs, WindowLimitedTransformer
-from fish_speech.models.dac.rvq import DownsampleResidualVectorQuantize
-
-
-def _rvq_transformer() -> WindowLimitedTransformer:
-    return WindowLimitedTransformer(
-        causal=True,
-        window_size=128,
-        input_dim=1024,
-        config=ModelArgs(
-            block_size=2048,
-            n_layer=8,
-            n_head=16,
-            dim=1024,
-            intermediate_size=3072,
-            n_local_heads=-1,
-            head_dim=64,
-            rope_base=10000,
-            norm_eps=1e-5,
-            dropout_rate=0.1,
-            attn_dropout_rate=0.1,
-            channels_first=True,
-        ),
-    )
-
-
-def build_codec() -> DAC:
-    quantizer = DownsampleResidualVectorQuantize(
-        input_dim=1024,
-        n_codebooks=9,
-        codebook_size=1024,
-        codebook_dim=8,
-        quantizer_dropout=0.5,
-        downsample_factor=[2, 2],
-        post_module=_rvq_transformer(),
-        pre_module=_rvq_transformer(),
-        semantic_codebook_size=4096,
-    )
-    return DAC(
-        sample_rate=44100,
-        encoder_dim=64,
-        encoder_rates=[2, 4, 8, 8],
-        decoder_dim=1536,
-        decoder_rates=[8, 8, 4, 2],
-        encoder_transformer_layers=[0, 0, 0, 4],
-        decoder_transformer_layers=[4, 0, 0, 0],
-        quantizer=quantizer,
-        transformer_general_config=partial(
-            ModelArgs,
-            block_size=8192,
-            n_local_heads=-1,
-            head_dim=64,
-            rope_base=10000,
-            norm_eps=1e-5,
-            dropout_rate=0.1,
-            attn_dropout_rate=0.1,
-            channels_first=True,
-        ),
-    )
-
-
-@torch.no_grad()
-def load_codec(path: str, dtype: torch.dtype = torch.bfloat16) -> DAC:
-    codec = build_codec()
-    state_dict = torch.load(path, map_location="cpu", mmap=False, weights_only=True)
-    if "state_dict" in state_dict:
-        state_dict = state_dict["state_dict"]
-    if any("generator" in k for k in state_dict):
-        state_dict = {
-            k.replace("generator.", ""): v
-            for k, v in state_dict.items()
-            if "generator." in k
-        }
-    rebuilt_buffers = {
-        f"{name}.{buffer}" if name else buffer
-        for name, module in codec.named_modules()
-        for buffer in module._non_persistent_buffers_set
-    }
-    state_dict = {k: v for k, v in state_dict.items() if k not in rebuilt_buffers}
-    codec.load_state_dict(state_dict, strict=True)
-    codec.eval()
-    for module in list(codec.modules()):
-        if is_parametrized(module, "weight"):
-            remove_parametrizations(module, "weight", leave_parametrized=True)
-        elif hasattr(module, "weight_g"):
-            torch.nn.utils.remove_weight_norm(module)
-    return codec.to(dtype=dtype)
-
-
-@torch.no_grad()
-def encode(codec: DAC, wav: torch.Tensor) -> torch.Tensor:
-    """Mono float waveform at codec.sample_rate -> codes [num_codebooks, T]."""
-    device = next(codec.parameters()).device
-    dtype = next(codec.parameters()).dtype
-    audios = wav.to(device=device, dtype=dtype)[None, None]
-    lengths = torch.tensor([wav.shape[-1]], device=device, dtype=torch.long)
-    indices, feature_lengths = codec.encode(audios, lengths)
-    return indices[0, :, : feature_lengths[0]]
-
-
-@torch.no_grad()
-def decode(codec: DAC, codes: torch.Tensor) -> torch.Tensor:
-    device = next(codec.parameters()).device
-    return codec.from_indices(codes.to(device)[None])[0, 0]

fish_speech/models/__init__.py

@@ -1 +0,0 @@
-

fish_speech/models/dac/__init__.py

@@ -1 +0,0 @@
-

fish_speech/models/dac/dac_nn.py

@@ -1,285 +0,0 @@
-# Vendored from descript-audio-codec (MIT License, Copyright (c) 2023-present Descript):
-# dac/nn/layers.py and dac/nn/quantize.py, with TorchScript removed from snake().
-from typing import Union
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange
-from torch.nn.utils import weight_norm
-
-
-def WNConv1d(*args, **kwargs):
-    return weight_norm(nn.Conv1d(*args, **kwargs))
-
-
-def WNConvTranspose1d(*args, **kwargs):
-    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
-
-
-# Scripting this brings model speed up 1.4x
-
-def snake(x, alpha):
-    shape = x.shape
-    x = x.reshape(shape[0], shape[1], -1)
-    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
-    x = x.reshape(shape)
-    return x
-
-
-class Snake1d(nn.Module):
-    def __init__(self, channels):
-        super().__init__()
-        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
-
-    def forward(self, x):
-        return snake(x, self.alpha)
-
-
-
-
-
-class VectorQuantize(nn.Module):
-    """
-    Implementation of VQ similar to Karpathy's repo:
-    https://github.com/karpathy/deep-vector-quantization
-    Additionally uses following tricks from Improved VQGAN
-    (https://arxiv.org/pdf/2110.04627.pdf):
-        1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space
-            for improved codebook usage
-        2. l2-normalized codes: Converts euclidean distance to cosine similarity which
-            improves training stability
-    """
-
-    def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int):
-        super().__init__()
-        self.codebook_size = codebook_size
-        self.codebook_dim = codebook_dim
-
-        self.in_proj = WNConv1d(input_dim, codebook_dim, kernel_size=1)
-        self.out_proj = WNConv1d(codebook_dim, input_dim, kernel_size=1)
-        self.codebook = nn.Embedding(codebook_size, codebook_dim)
-
-    def forward(self, z):
-        """Quantized the input tensor using a fixed codebook and returns
-        the corresponding codebook vectors
-
-        Parameters
-        ----------
-        z : Tensor[B x D x T]
-
-        Returns
-        -------
-        Tensor[B x D x T]
-            Quantized continuous representation of input
-        Tensor[1]
-            Commitment loss to train encoder to predict vectors closer to codebook
-            entries
-        Tensor[1]
-            Codebook loss to update the codebook
-        Tensor[B x T]
-            Codebook indices (quantized discrete representation of input)
-        Tensor[B x D x T]
-            Projected latents (continuous representation of input before quantization)
-        """
-
-        # Factorized codes (ViT-VQGAN) Project input into low-dimensional space
-        z_e = self.in_proj(z)  # z_e : (B x D x T)
-        z_q, indices = self.decode_latents(z_e)
-
-        commitment_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2])
-        codebook_loss = F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2])
-
-        z_q = (
-            z_e + (z_q - z_e).detach()
-        )  # noop in forward pass, straight-through gradient estimator in backward pass
-
-        z_q = self.out_proj(z_q)
-
-        return z_q, commitment_loss, codebook_loss, indices, z_e
-
-    def embed_code(self, embed_id):
-        return F.embedding(embed_id, self.codebook.weight)
-
-    def decode_code(self, embed_id):
-        return self.embed_code(embed_id).transpose(1, 2)
-
-    def decode_latents(self, latents):
-        encodings = rearrange(latents, "b d t -> (b t) d")
-        codebook = self.codebook.weight  # codebook: (N x D)
-
-        # L2 normalize encodings and codebook (ViT-VQGAN)
-        encodings = F.normalize(encodings)
-        codebook = F.normalize(codebook)
-
-        # Compute euclidean distance with codebook
-        dist = (
-            encodings.pow(2).sum(1, keepdim=True)
-            - 2 * encodings @ codebook.t()
-            + codebook.pow(2).sum(1, keepdim=True).t()
-        )
-        indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
-        z_q = self.decode_code(indices)
-        return z_q, indices
-
-
-class ResidualVectorQuantize(nn.Module):
-    """
-    Introduced in SoundStream: An end2end neural audio codec
-    https://arxiv.org/abs/2107.03312
-    """
-
-    def __init__(
-        self,
-        input_dim: int = 512,
-        n_codebooks: int = 9,
-        codebook_size: int = 1024,
-        codebook_dim: Union[int, list] = 8,
-        quantizer_dropout: float = 0.0,
-    ):
-        super().__init__()
-        if isinstance(codebook_dim, int):
-            codebook_dim = [codebook_dim for _ in range(n_codebooks)]
-
-        self.n_codebooks = n_codebooks
-        self.codebook_dim = codebook_dim
-        self.codebook_size = codebook_size
-
-        self.quantizers = nn.ModuleList(
-            [
-                VectorQuantize(input_dim, codebook_size, codebook_dim[i])
-                for i in range(n_codebooks)
-            ]
-        )
-        self.quantizer_dropout = quantizer_dropout
-
-    def forward(self, z, n_quantizers: int = None):
-        """Quantized the input tensor using a fixed set of `n` codebooks and returns
-        the corresponding codebook vectors
-        Parameters
-        ----------
-        z : Tensor[B x D x T]
-        n_quantizers : int, optional
-            No. of quantizers to use
-            (n_quantizers < self.n_codebooks ex: for quantizer dropout)
-            Note: if `self.quantizer_dropout` is True, this argument is ignored
-                when in training mode, and a random number of quantizers is used.
-        Returns
-        -------
-        dict
-            A dictionary with the following keys:
-
-            "z" : Tensor[B x D x T]
-                Quantized continuous representation of input
-            "codes" : Tensor[B x N x T]
-                Codebook indices for each codebook
-                (quantized discrete representation of input)
-            "latents" : Tensor[B x N*D x T]
-                Projected latents (continuous representation of input before quantization)
-            "vq/commitment_loss" : Tensor[1]
-                Commitment loss to train encoder to predict vectors closer to codebook
-                entries
-            "vq/codebook_loss" : Tensor[1]
-                Codebook loss to update the codebook
-        """
-        z_q = 0
-        residual = z
-        commitment_loss = 0
-        codebook_loss = 0
-
-        codebook_indices = []
-        latents = []
-
-        if n_quantizers is None:
-            n_quantizers = self.n_codebooks
-        if self.training:
-            n_quantizers = torch.ones((z.shape[0],)) * self.n_codebooks + 1
-            dropout = torch.randint(1, self.n_codebooks + 1, (z.shape[0],))
-            n_dropout = int(z.shape[0] * self.quantizer_dropout)
-            n_quantizers[:n_dropout] = dropout[:n_dropout]
-            n_quantizers = n_quantizers.to(z.device)
-
-        for i, quantizer in enumerate(self.quantizers):
-            if self.training is False and i >= n_quantizers:
-                break
-
-            z_q_i, commitment_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer(
-                residual
-            )
-
-            # Create mask to apply quantizer dropout
-            mask = (
-                torch.full((z.shape[0],), fill_value=i, device=z.device) < n_quantizers
-            )
-            z_q = z_q + z_q_i * mask[:, None, None]
-            residual = residual - z_q_i
-
-            # Sum losses
-            commitment_loss += (commitment_loss_i * mask).mean()
-            codebook_loss += (codebook_loss_i * mask).mean()
-
-            codebook_indices.append(indices_i)
-            latents.append(z_e_i)
-
-        codes = torch.stack(codebook_indices, dim=1)
-        latents = torch.cat(latents, dim=1)
-
-        return z_q, codes, latents, commitment_loss, codebook_loss
-
-    def from_codes(self, codes: torch.Tensor):
-        """Given the quantized codes, reconstruct the continuous representation
-        Parameters
-        ----------
-        codes : Tensor[B x N x T]
-            Quantized discrete representation of input
-        Returns
-        -------
-        Tensor[B x D x T]
-            Quantized continuous representation of input
-        """
-        z_q = 0.0
-        z_p = []
-        n_codebooks = codes.shape[1]
-        for i in range(n_codebooks):
-            z_p_i = self.quantizers[i].decode_code(codes[:, i, :])
-            z_p.append(z_p_i)
-
-            z_q_i = self.quantizers[i].out_proj(z_p_i)
-            z_q = z_q + z_q_i
-        return z_q, torch.cat(z_p, dim=1), codes
-
-    def from_latents(self, latents: torch.Tensor):
-        """Given the unquantized latents, reconstruct the
-        continuous representation after quantization.
-
-        Parameters
-        ----------
-        latents : Tensor[B x N x T]
-            Continuous representation of input after projection
-
-        Returns
-        -------
-        Tensor[B x D x T]
-            Quantized representation of full-projected space
-        Tensor[B x D x T]
-            Quantized representation of latent space
-        """
-        z_q = 0
-        z_p = []
-        codes = []
-        dims = np.cumsum([0] + [q.codebook_dim for q in self.quantizers])
-
-        n_codebooks = np.where(dims <= latents.shape[1])[0].max(axis=0, keepdims=True)[
-            0
-        ]
-        for i in range(n_codebooks):
-            j, k = dims[i], dims[i + 1]
-            z_p_i, codes_i = self.quantizers[i].decode_latents(latents[:, j:k, :])
-            z_p.append(z_p_i)
-            codes.append(codes_i)
-
-            z_q_i = self.quantizers[i].out_proj(z_p_i)
-            z_q = z_q + z_q_i
-
-        return z_q, torch.cat(z_p, dim=1), torch.stack(codes, dim=1)

fish_speech/models/dac/modded_dac.py

@@ -1,991 +0,0 @@
-import math
-import typing as tp
-from dataclasses import dataclass
-from typing import List, Optional, Union
-
-import numpy as np
-import torch
-from fish_speech.models.dac.dac_nn import Snake1d, WNConv1d, WNConvTranspose1d
-from torch import Tensor, nn
-from torch.nn import functional as F
-from torch.nn.utils.parametrizations import weight_norm
-from torch.nn.utils.parametrize import remove_parametrizations
-
-
-@dataclass
-class VQResult:
-    z: torch.Tensor
-    codes: torch.Tensor
-    latents: torch.Tensor
-    codebook_loss: torch.Tensor
-    commitment_loss: torch.Tensor
-    semantic_distill_z: torch.Tensor | None = None
-
-
-def find_multiple(n: int, k: int) -> int:
-    if n % k == 0:
-        return n
-    return n + k - (n % k)
-
-
-@dataclass
-class ModelArgs:
-    block_size: int = 2048
-    n_layer: int = 8
-    n_head: int = 8
-    dim: int = 512
-    intermediate_size: int = 1536
-    n_local_heads: int = -1
-    head_dim: int = 64
-    rope_base: float = 10000
-    norm_eps: float = 1e-5
-    dropout_rate: float = 0.1
-    attn_dropout_rate: float = 0.1
-    channels_first: bool = True  # to be compatible with conv1d input/output
-    pos_embed_type: str = "rope"  # can be "rope" or "conformer"
-    max_relative_position: int = 128  # for conformer-style relative position embedding
-    window_size: int = 512  # for window limited attention
-
-    def __post_init__(self):
-        if self.n_local_heads == -1:
-            self.n_local_heads = self.n_head
-        if self.intermediate_size is None:
-            hidden_dim = 4 * self.dim
-            n_hidden = int(2 * hidden_dim / 3)
-            self.intermediate_size = find_multiple(n_hidden, 256)
-        assert self.pos_embed_type in [
-            "rope",
-            "conformer",
-        ], "pos_embed_type must be either 'rope' or 'conformer'"
-
-
-class KVCache(nn.Module):
-    def __init__(
-        self, max_batch_size, max_seq_length, n_heads, head_dim, dtype=torch.bfloat16
-    ):
-        super().__init__()
-        cache_shape = (max_batch_size, n_heads, max_seq_length, head_dim)
-        self.register_buffer("k_cache", torch.zeros(cache_shape, dtype=dtype))
-        self.register_buffer("v_cache", torch.zeros(cache_shape, dtype=dtype))
-
-    def update(self, input_pos, k_val, v_val):
-        # input_pos: [S], k_val: [B, H, S, D]
-        assert input_pos.shape[0] == k_val.shape[2]
-
-        k_out = self.k_cache
-        v_out = self.v_cache
-        k_out[:, :, input_pos] = k_val
-        v_out[:, :, input_pos] = v_val
-
-        return (
-            k_out[:, :, : input_pos.max() + 1, :],
-            v_out[:, :, : input_pos.max() + 1, :],
-        )
-
-    def clear_cache(self, prompt_len):
-        self.k_cache[:, :, prompt_len:, :] = torch.zeros_like(
-            self.k_cache[:, :, prompt_len:, :]
-        )
-        self.v_cache[:, :, prompt_len:, :] = torch.zeros_like(
-            self.v_cache[:, :, prompt_len:, :]
-        )
-
-
-class Transformer(nn.Module):
-    def __init__(self, config: ModelArgs) -> None:
-        super().__init__()
-        self.config = config
-
-        self.layers = nn.ModuleList(
-            TransformerBlock(config) for _ in range(config.n_layer)
-        )
-        self.norm = RMSNorm(config.dim, eps=config.norm_eps)
-
-        # Only compute RoPE frequencies if using RoPE
-        if config.pos_embed_type == "rope":
-            freqs_cis = precompute_freqs_cis(
-                327680, self.config.head_dim, self.config.rope_base
-            )
-            self.register_buffer("freqs_cis", freqs_cis, persistent=False)
-        else:
-            self.register_buffer("freqs_cis", None)
-
-        causal_mask = torch.tril(torch.ones(32768, 32768, dtype=torch.bool))
-        self.register_buffer("causal_mask", causal_mask, persistent=False)
-
-        self.max_batch_size = -1
-        self.max_seq_length = -1
-        self.use_kv_cache = False
-
-    def setup_caches(self, max_batch_size, max_seq_length):
-        """
-        This method will only be called during inference when using KV cache.
-        """
-        head_dim = self.config.dim // self.config.n_head
-        max_seq_length = find_multiple(max_seq_length, 8)
-        self.max_seq_length = max_seq_length
-        self.max_batch_size = max_batch_size
-        dtype = self.norm.weight.dtype
-        device = self.norm.weight.device
-
-        for b in self.layers:
-            b.attention.kv_cache = KVCache(
-                max_batch_size,
-                max_seq_length,
-                self.config.n_local_heads,
-                head_dim,
-                dtype,
-            ).to(device)
-
-        self.use_kv_cache = True
-
-    def forward(
-        self,
-        x: Tensor,
-        input_pos: Optional[Tensor] = None,
-        mask: Optional[Tensor] = None,
-    ) -> Tensor:
-        if self.config.pos_embed_type == "rope":
-            assert (
-                self.freqs_cis is not None
-            ), "RoPE frequencies must be initialized for RoPE positional embedding"
-            # print("MAX", input_pos.max())
-            freqs_cis = self.freqs_cis[input_pos]
-        else:
-            freqs_cis = None
-
-        if mask is None:  # in case of non-causal model
-            if not self.training and self.use_kv_cache:
-                mask = self.causal_mask[None, None, input_pos]
-                mask = mask[..., : input_pos.max() + 1]
-            else:
-                mask = self.causal_mask[None, None, input_pos]
-                mask = mask[..., input_pos]
-
-        for i, layer in enumerate(self.layers):
-            x = layer(x, input_pos, freqs_cis, mask)
-        x = self.norm(x)
-        return x
-
-
-class TransformerBlock(nn.Module):
-    def __init__(self, config: ModelArgs) -> None:
-        super().__init__()
-        self.attention = Attention(config)
-        self.feed_forward = FeedForward(config)
-        self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)
-        self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)
-        self.attention_layer_scale = LayerScale(config.dim, inplace=True)
-        self.ffn_layer_scale = LayerScale(config.dim, inplace=True)
-
-    def forward(
-        self,
-        x: Tensor,
-        input_pos: Tensor,
-        freqs_cis: Tensor,
-        mask: Tensor,
-    ) -> Tensor:
-        h = x + self.attention_layer_scale(
-            self.attention(self.attention_norm(x), freqs_cis, mask, input_pos)
-        )
-        out = h + self.ffn_layer_scale(self.feed_forward(self.ffn_norm(h)))
-        return out
-
-
-class Attention(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        assert config.dim % config.n_head == 0
-
-        total_head_dim = (config.n_head + 2 * config.n_local_heads) * config.head_dim
-        # key, query, value projections for all heads, but in a batch
-        self.wqkv = nn.Linear(config.dim, total_head_dim, bias=False)
-        self.wo = nn.Linear(config.head_dim * config.n_head, config.dim, bias=False)
-        self.kv_cache = None
-
-        self.n_head = config.n_head
-        self.head_dim = config.head_dim
-        self.n_local_heads = config.n_local_heads
-        self.dim = config.dim
-        self.attn_dropout_rate = config.attn_dropout_rate
-        self.pos_embed_type = config.pos_embed_type
-
-        # Add relative position embedding for conformer-style
-        if self.pos_embed_type == "conformer":
-            self.max_relative_position = config.max_relative_position
-            num_pos_embeddings = 2 * config.max_relative_position + 1
-            self.rel_pos_embeddings = nn.Parameter(
-                torch.zeros(num_pos_embeddings, self.head_dim)
-            )
-            nn.init.normal_(self.rel_pos_embeddings, mean=0.0, std=0.02)
-
-    def _compute_conformer_pos_scores(self, q: Tensor, seqlen: int) -> Tensor:
-        # q: [B, H, S, D]
-        # Returns: [B, H, S, S]
-        positions = torch.arange(seqlen, device=q.device)
-        relative_positions = positions.unsqueeze(1) - positions.unsqueeze(0)  # [S, S]
-        relative_positions = torch.clamp(
-            relative_positions + self.max_relative_position,
-            0,
-            2 * self.max_relative_position,
-        )
-        rel_embeddings = self.rel_pos_embeddings[relative_positions]  # [S, S, D]
-
-        # Compute attention scores with relative position embeddings
-        q = q.transpose(1, 2)  # [B, S, H, D]
-        rel_logits = torch.matmul(q, rel_embeddings.transpose(-2, -1))  # [B, S, H, S]
-        rel_logits = rel_logits.transpose(1, 2)  # [B, H, S, S]
-        return rel_logits
-
-    def forward(
-        self,
-        x: Tensor,
-        freqs_cis: Tensor,
-        mask: Tensor,
-        input_pos: Optional[Tensor] = None,
-    ) -> Tensor:
-        bsz, seqlen, _ = x.shape
-
-        kv_size = self.n_local_heads * self.head_dim
-        q, k, v = self.wqkv(x).split([kv_size, kv_size, kv_size], dim=-1)
-        context_seqlen = seqlen
-
-        q = q.view(bsz, seqlen, self.n_head, self.head_dim)
-        k = k.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)
-        v = v.view(bsz, context_seqlen, self.n_local_heads, self.head_dim)
-
-        if self.pos_embed_type == "rope":
-            q = apply_rotary_emb(q, freqs_cis)
-            k = apply_rotary_emb(k, freqs_cis)
-
-        q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v))
-
-        if self.kv_cache is not None:
-            k, v = self.kv_cache.update(input_pos, k, v)
-
-        k = k.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
-        v = v.repeat_interleave(self.n_head // self.n_local_heads, dim=1)
-
-        if self.pos_embed_type == "conformer":
-            # Compute attention scores
-            scale = 1.0 / math.sqrt(self.head_dim)
-            scores = torch.matmul(q, k.transpose(-2, -1)) * scale
-
-            # Add relative position embeddings for conformer-style
-            rel_scores = self._compute_conformer_pos_scores(q, seqlen)
-            scores = scores + rel_scores
-
-            # Apply attention
-            if mask is not None:
-                scores = scores.masked_fill(~mask, float("-inf"))
-
-            attn = F.softmax(scores, dim=-1)
-            if self.attn_dropout_rate > 0 and self.training:
-                attn = F.dropout(attn, p=self.attn_dropout_rate)
-
-            y = torch.matmul(attn, v)
-        else:
-            y = F.scaled_dot_product_attention(
-                q,
-                k,
-                v,
-                dropout_p=self.attn_dropout_rate if self.training else 0.0,
-                attn_mask=mask,
-            )
-            # is_causal=True)
-        y = (
-            y.transpose(1, 2)
-            .contiguous()
-            .view(bsz, seqlen, self.head_dim * self.n_head)
-        )
-        y = self.wo(y)
-        return y
-
-
-class FeedForward(nn.Module):
-    def __init__(self, config: ModelArgs) -> None:
-        super().__init__()
-        self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False)
-        self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False)
-        self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False)
-        self.dropout = nn.Dropout(config.dropout_rate)
-
-    def forward(self, x: Tensor) -> Tensor:
-        return self.w2(self.dropout(F.silu(self.w1(x)) * self.w3(x)))
-
-
-class RMSNorm(nn.Module):
-    def __init__(self, dim: int, eps: float = 1e-5):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim))
-
-    def _norm(self, x):
-        return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
-
-    def forward(self, x: Tensor) -> Tensor:
-        output = self._norm(x.float()).type_as(x)
-        return output * self.weight
-
-
-class LayerScale(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        init_values: Union[float, Tensor] = 1e-2,
-        inplace: bool = False,
-    ) -> None:
-        super().__init__()
-        self.inplace = inplace
-        self.gamma = nn.Parameter(init_values * torch.ones(dim))
-
-    def forward(self, x: Tensor) -> Tensor:
-        return x.mul_(self.gamma) if self.inplace else x * self.gamma
-
-
-class WindowLimitedTransformer(Transformer):
-    """
-    Transformer with window limited attention, causal.
-    """
-
-    def __init__(
-        self,
-        config: ModelArgs,
-        input_dim: int = 512,
-        window_size: Optional[int] = None,
-        causal: bool = True,
-        look_ahead_conv: nn.Module = None,
-    ):
-        super().__init__(config)
-        self.window_size = window_size
-        self.causal = causal
-        self.channels_first = config.channels_first
-        self.look_ahead_conv = (
-            look_ahead_conv if look_ahead_conv is not None else nn.Identity()
-        )
-        self.input_proj = (
-            nn.Linear(input_dim, config.dim)
-            if input_dim != config.dim
-            else nn.Identity()
-        )
-        self.output_proj = (
-            nn.Linear(config.dim, input_dim)
-            if input_dim != config.dim
-            else nn.Identity()
-        )
-
-    def make_window_limited_mask(
-        self,
-        max_length: int,
-        x_lens: Optional[Tensor] = None,
-    ) -> Tensor:
-        """
-        Make mask to form window limited attention.
-        """
-        if self.causal:
-            mask = torch.tril(torch.ones(max_length, max_length))
-            row_indices = torch.arange(max_length).view(-1, 1)
-            window_size = self.window_size or max_length
-            valid_range = (row_indices - window_size + 1).clamp(min=0)
-            column_indices = torch.arange(max_length)
-            mask = (column_indices >= valid_range) & mask.bool()
-        else:
-            raise NotImplementedError
-        mask = mask.bool()[None, None]
-        return mask
-
-    def make_mask(
-        self,
-        max_length: int,
-        x_lens: Optional[Tensor] = None,
-    ) -> Tensor:
-        """
-        Make ordinary mask if window size is not specified.
-        """
-        if self.causal:
-            mask = torch.tril(torch.ones(max_length, max_length))
-        else:
-            mask = torch.ones(max_length, max_length)
-            mask = mask.bool()[None, None]
-            for i, x_len in enumerate(x_lens):
-                mask[:x_len, i] = 0
-        mask = mask.bool()[None, None]
-        return mask
-
-    def forward(
-        self,
-        x: Tensor,
-        x_lens: Optional[Tensor] = None,
-    ) -> Tensor:
-        if self.channels_first:
-            x = x.transpose(1, 2)
-        x = self.input_proj(x)  # (B, T, D)
-        x = self.look_ahead_conv(x)
-        input_pos = torch.arange(x.shape[1], device=x.device)
-        # construct mask to form window limited attention
-        max_length = x.shape[1]
-        if self.window_size is not None:
-            mask = self.make_window_limited_mask(max_length, x_lens)
-        else:
-            mask = self.make_mask(max_length, x_lens)
-        mask = mask.to(x.device)
-        x = super().forward(x, input_pos, mask)
-        x = self.output_proj(x)  # (B, T, D)
-        if self.channels_first:
-            x = x.transpose(1, 2)
-        return x
-
-
-def precompute_freqs_cis(
-    seq_len: int, n_elem: int, base: int = 10000, dtype: torch.dtype = torch.bfloat16
-) -> Tensor:
-    freqs = 1.0 / (
-        base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem)
-    )
-    t = torch.arange(seq_len, device=freqs.device)
-    freqs = torch.outer(t, freqs)
-    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
-    cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1)
-    return cache.to(dtype=dtype)
-
-
-def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
-    xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
-    freqs_cis = freqs_cis.view(1, xshaped.size(1), 1, xshaped.size(3), 2)
-    x_out2 = torch.stack(
-        [
-            xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1],
-            xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1],
-        ],
-        -1,
-    )
-
-    x_out2 = x_out2.flatten(3)
-    return x_out2.type_as(x)
-
-
-def init_weights(m):
-    if isinstance(m, nn.Conv1d):
-        nn.init.trunc_normal_(m.weight, std=0.02)
-        nn.init.constant_(m.bias, 0)
-
-
-def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]):
-    """Remove padding from x, handling properly zero padding. Only for 1d!"""
-    padding_left, padding_right = paddings
-    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
-    assert (padding_left + padding_right) <= x.shape[-1]
-    end = x.shape[-1] - padding_right
-    return x[..., padding_left:end]
-
-
-def get_extra_padding_for_conv1d(
-    x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0
-) -> int:
-    """See `pad_for_conv1d`."""
-    length = x.shape[-1]
-    n_frames = (length - kernel_size + padding_total) / stride + 1
-    ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
-    return ideal_length - length
-
-
-def pad1d(
-    x: torch.Tensor,
-    paddings: tp.Tuple[int, int],
-    mode: str = "zeros",
-    value: float = 0.0,
-):
-    """Tiny wrapper around F.pad, just to allow for reflect padding on small input.
-    If this is the case, we insert extra 0 padding to the right
-    before the reflection happen.
-    """
-    length = x.shape[-1]
-    padding_left, padding_right = paddings
-    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
-    if mode == "reflect":
-        max_pad = max(padding_left, padding_right)
-        extra_pad = 0
-        if length <= max_pad:
-            extra_pad = max_pad - length + 1
-            x = F.pad(x, (0, extra_pad))
-        padded = F.pad(x, paddings, mode, value)
-        end = padded.shape[-1] - extra_pad
-        return padded[..., :end]
-    else:
-        return F.pad(x, paddings, mode, value)
-
-
-class CausalConvNet(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        kernel_size,
-        dilation=1,
-        stride=1,
-        groups=1,
-        padding=None,
-    ):
-        super(CausalConvNet, self).__init__()
-        self.conv = nn.Conv1d(
-            in_channels,
-            out_channels,
-            kernel_size,
-            stride=stride,
-            dilation=dilation,
-            groups=groups,
-        )
-        self.stride = stride
-        self.kernel_size = (kernel_size - 1) * dilation + 1
-        self.dilation = dilation
-        self.padding = self.kernel_size - self.stride
-
-    def forward(self, x):
-        pad = self.padding
-        extra_padding = get_extra_padding_for_conv1d(
-            x, self.kernel_size, self.stride, pad
-        )
-        x = pad1d(x, (pad, extra_padding), mode="constant", value=0)
-        return self.conv(x).contiguous()
-
-    def weight_norm(self, name="weight", dim=0):
-        self.conv = weight_norm(self.conv, name=name, dim=dim)
-        return self
-
-    def remove_weight_norm(self):
-        self.conv = remove_parametrizations(self.conv)
-        return self
-
-
-class CausalTransConvNet(nn.Module):
-    def __init__(
-        self, in_channels, out_channels, kernel_size, dilation=1, stride=1, padding=None
-    ):
-        super(CausalTransConvNet, self).__init__()
-        self.conv = nn.ConvTranspose1d(
-            in_channels, out_channels, kernel_size, stride=stride, dilation=dilation
-        )
-        self.stride = stride
-        self.kernel_size = kernel_size
-
-    def forward(self, x):
-        x = self.conv(x)
-        pad = self.kernel_size - self.stride
-        padding_right = math.ceil(pad)
-        padding_left = pad - padding_right
-        x = unpad1d(x, (padding_left, padding_right))
-        return x.contiguous()
-
-    def weight_norm(self, name="weight", dim=0):
-        self.conv = weight_norm(self.conv, name=name, dim=dim)
-        return self
-
-    def remove_weight_norm(self):
-        self.conv = remove_parametrizations(self.conv)
-        return self
-
-
-def CausalWNConv1d(*args, **kwargs):
-    return CausalConvNet(*args, **kwargs).weight_norm()
-
-
-def CausalWNConvTranspose1d(*args, **kwargs):
-    return CausalTransConvNet(*args, **kwargs).weight_norm()
-
-
-class ResidualUnit(nn.Module):
-    def __init__(self, dim: int = 16, dilation: int = 1, causal: bool = False):
-        super().__init__()
-        conv_class = CausalWNConv1d if causal else WNConv1d
-        pad = ((7 - 1) * dilation) // 2
-        self.block = nn.Sequential(
-            Snake1d(dim),
-            conv_class(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
-            Snake1d(dim),
-            conv_class(dim, dim, kernel_size=1),
-        )
-        self.causal = causal
-
-    def forward(self, x):
-        y = self.block(x)
-        pad = x.shape[-1] - y.shape[-1]
-        if pad > 0:
-            if self.causal:
-                x = x[..., :-pad]
-            else:
-                x = x[..., pad // 2 : -pad // 2]
-        return x + y
-
-
-class EncoderBlock(nn.Module):
-    def __init__(
-        self,
-        dim: int = 16,
-        stride: int = 1,
-        causal: bool = False,
-        n_t_layer: int = 0,
-        transformer_general_config=None,
-    ):
-        super().__init__()
-        conv_class = CausalWNConv1d if causal else WNConv1d
-        transformer_module = (
-            nn.Identity()
-            if n_t_layer == 0
-            else (
-                WindowLimitedTransformer(
-                    causal=causal,
-                    input_dim=dim,
-                    window_size=getattr(transformer_general_config, "window_size", 512),
-                    config=transformer_general_config(
-                        n_layer=n_t_layer,
-                        n_head=dim // 64,
-                        dim=dim,
-                        intermediate_size=dim * 3,
-                    ),
-                )
-            )
-        )
-        self.block = nn.Sequential(
-            ResidualUnit(dim // 2, dilation=1, causal=causal),
-            ResidualUnit(dim // 2, dilation=3, causal=causal),
-            ResidualUnit(dim // 2, dilation=9, causal=causal),
-            Snake1d(dim // 2),
-            conv_class(
-                dim // 2,
-                dim,
-                kernel_size=2 * stride,
-                stride=stride,
-                padding=math.ceil(stride / 2),
-            ),
-            transformer_module,
-        )
-
-    def forward(self, x):
-        return self.block(x)
-
-
-class Encoder(nn.Module):
-    def __init__(
-        self,
-        d_model: int = 64,
-        strides: list = [2, 4, 8, 8],
-        d_latent: int = 64,
-        n_transformer_layers: list = [0, 0, 4, 4],
-        transformer_general_config: ModelArgs = None,
-        causal: bool = False,
-    ):
-        super().__init__()
-        conv_class = CausalWNConv1d if causal else WNConv1d
-        # Create first convolution
-        self.block = [conv_class(1, d_model, kernel_size=7, padding=3)]
-
-        # Create EncoderBlocks that double channels as they downsample by `stride`
-        for stride, n_t_layer in zip(strides, n_transformer_layers):
-            d_model *= 2
-            self.block += [
-                EncoderBlock(
-                    d_model,
-                    stride=stride,
-                    causal=causal,
-                    n_t_layer=n_t_layer,
-                    transformer_general_config=transformer_general_config,
-                )
-            ]
-
-        # Create last convolution
-        self.block += [
-            Snake1d(d_model),
-            conv_class(d_model, d_latent, kernel_size=3, padding=1),
-        ]
-
-        # Wrap black into nn.Sequential
-        self.block = nn.Sequential(*self.block)
-        self.enc_dim = d_model
-
-    def forward(self, x):
-        return self.block(x)
-
-
-class DecoderBlock(nn.Module):
-    def __init__(
-        self,
-        input_dim: int = 16,
-        output_dim: int = 8,
-        stride: int = 1,
-        causal: bool = False,
-        n_t_layer: int = 0,
-        transformer_general_config=None,
-    ):
-        super().__init__()
-        conv_trans_class = CausalWNConvTranspose1d if causal else WNConvTranspose1d
-        transformer_module = (
-            nn.Identity()
-            if n_t_layer == 0
-            else (
-                WindowLimitedTransformer(
-                    causal=causal,
-                    input_dim=input_dim,
-                    window_size=None,
-                    config=transformer_general_config(
-                        n_layer=n_t_layer,
-                        n_head=input_dim // 64,
-                        dim=input_dim,
-                        intermediate_size=input_dim * 3,
-                    ),
-                )
-            )
-        )
-        self.block = nn.Sequential(
-            # transformer_module,
-            Snake1d(input_dim),
-            conv_trans_class(
-                input_dim,
-                output_dim,
-                kernel_size=2 * stride,
-                stride=stride,
-                padding=math.ceil(stride / 2),
-            ),
-            ResidualUnit(output_dim, dilation=1, causal=causal),
-            ResidualUnit(output_dim, dilation=3, causal=causal),
-            ResidualUnit(output_dim, dilation=9, causal=causal),
-        )
-
-    def forward(self, x):
-        return self.block(x)
-
-
-class Decoder(nn.Module):
-    def __init__(
-        self,
-        input_channel,
-        channels,
-        rates,
-        d_out: int = 1,
-        causal: bool = False,
-        n_transformer_layers: list = [0, 0, 0, 0],
-        transformer_general_config=None,
-    ):
-        super().__init__()
-        conv_class = CausalWNConv1d if causal else WNConv1d
-        # Add first conv layer
-        layers = [conv_class(input_channel, channels, kernel_size=7, padding=3)]
-
-        # Add upsampling + MRF blocks
-        for i, (stride, n_t_layer) in enumerate(zip(rates, n_transformer_layers)):
-            input_dim = channels // 2**i
-            output_dim = channels // 2 ** (i + 1)
-            layers += [
-                DecoderBlock(
-                    input_dim,
-                    output_dim,
-                    stride,
-                    causal=causal,
-                    n_t_layer=n_t_layer,
-                    transformer_general_config=transformer_general_config,
-                )
-            ]
-
-        # Add final conv layer
-        layers += [
-            Snake1d(output_dim),
-            conv_class(output_dim, d_out, kernel_size=7, padding=3),
-            nn.Tanh(),
-        ]
-
-        self.model = nn.Sequential(*layers)
-
-    def forward(self, x):
-        return self.model(x)
-
-
-class DAC(nn.Module):
-    def __init__(
-        self,
-        encoder_dim: int = 64,
-        encoder_rates: List[int] = [2, 4, 8, 8],
-        latent_dim: int = None,
-        decoder_dim: int = 1536,
-        decoder_rates: List[int] = [8, 8, 4, 2],
-        quantizer: torch.nn.Module = None,
-        sample_rate: int = 44100,
-        causal: bool = True,
-        encoder_transformer_layers: List[int] = [0, 0, 0, 0],
-        decoder_transformer_layers: List[int] = [0, 0, 0, 0],
-        overwrite_decoder: torch.nn.Module = None,
-        transformer_general_config=None,
-    ):
-        super().__init__()
-
-        self.encoder_dim = encoder_dim
-        self.encoder_rates = encoder_rates
-        self.decoder_dim = decoder_dim
-        self.decoder_rates = decoder_rates
-        self.sample_rate = sample_rate
-
-        if latent_dim is None:
-            latent_dim = encoder_dim * (2 ** len(encoder_rates))
-
-        self.latent_dim = latent_dim
-
-        self.hop_length = np.prod(encoder_rates)
-        self.encoder = Encoder(
-            encoder_dim,
-            encoder_rates,
-            latent_dim,
-            causal=causal,
-            n_transformer_layers=encoder_transformer_layers,
-            transformer_general_config=transformer_general_config,
-        )
-
-        self.quantizer = quantizer
-
-        if overwrite_decoder is not None:
-            self.decoder = overwrite_decoder
-        else:
-            self.decoder = Decoder(
-                latent_dim,
-                decoder_dim,
-                decoder_rates,
-                causal=causal,
-                n_transformer_layers=decoder_transformer_layers,
-                transformer_general_config=transformer_general_config,
-            )
-        self.sample_rate = sample_rate
-        self.apply(init_weights)
-
-        self.frame_length = self.hop_length * 4
-
-    def preprocess(self, audio_data, sample_rate):
-        if sample_rate is None:
-            sample_rate = self.sample_rate
-        assert sample_rate == self.sample_rate
-
-        length = audio_data.shape[-1]
-        right_pad = math.ceil(length / self.hop_length) * self.hop_length - length
-        audio_data = nn.functional.pad(audio_data, (0, right_pad))
-
-        return audio_data
-
-    def encode(
-        self,
-        audio_data: torch.Tensor,
-        audio_lengths: torch.Tensor = None,
-        n_quantizers: int = None,
-        **kwargs,
-    ):
-        """Encode given audio data and return quantized latent codes
-
-        Parameters
-        ----------
-        audio_data : Tensor[B x T]
-            Audio data to encode
-        n_quantizers : int, optional
-            Number of quantizers to use, by default None
-            If None, all quantizers are used.
-
-        Returns
-        -------
-        dict
-            A dictionary with the following keys:
-            "z" : Tensor[B x D x T]
-                Quantized continuous representation of input
-            "codes" : Tensor[B x N x T]
-                Codebook indices for each codebook
-                (quantized discrete representation of input)
-            "latents" : Tensor[B x N*D x T]
-                Projected latents (continuous representation of input before quantization)
-            "vq/commitment_loss" : Tensor[1]
-                Commitment loss to train encoder to predict vectors closer to codebook
-                entries
-            "vq/codebook_loss" : Tensor[1]
-                Codebook loss to update the codebook
-            "length" : int
-                Number of samples in input audio
-        """
-        # pad to multiple of self.frame_length
-        if audio_data.ndim == 2:
-            audio_data = audio_data.unsqueeze(1)
-        length = audio_data.shape[-1]
-        right_pad = math.ceil(length / self.frame_length) * self.frame_length - length
-        audio_data = nn.functional.pad(audio_data, (0, right_pad))
-        if audio_lengths is None:
-            audio_lengths = torch.LongTensor([length + right_pad]).to(audio_data.device)
-
-        z = self.encoder(audio_data)
-        vq_results = self.quantizer(z, n_quantizers, **kwargs)
-        indices = vq_results.codes
-        indices_lens = torch.ceil(audio_lengths / self.frame_length).long()
-        return indices, indices_lens
-
-    def from_indices(self, indices: torch.Tensor):
-        z = self.quantizer.decode(indices)
-        return self.decoder(z)
-
-    def decode(self, z: torch.Tensor):
-        """Decode given latent codes and return audio data
-
-        Parameters
-        ----------
-        z : Tensor[B x D x T]
-            Quantized continuous representation of input
-        length : int, optional
-            Number of samples in output audio, by default None
-
-        Returns
-        -------
-        dict
-            A dictionary with the following keys:
-            "audio" : Tensor[B x 1 x length]
-                Decoded audio data.
-        """
-        return self.decoder(z)
-
-    def forward(
-        self,
-        audio_data: torch.Tensor,
-        template: torch.Tensor = None,
-        mask: torch.Tensor = None,
-        sample_rate: int = None,
-        n_quantizers: int = None,
-        **kwargs,
-    ):
-        """Model forward pass
-
-        Parameters
-        ----------
-        audio_data : Tensor[B x 1 x T]
-            Audio data to encode
-        sample_rate : int, optional
-            Sample rate of audio data in Hz, by default None
-            If None, defaults to `self.sample_rate`
-        n_quantizers : int, optional
-            Number of quantizers to use, by default None.
-            If None, all quantizers are used.
-
-        Returns
-        -------
-        dict
-            A dictionary with the following keys:
-            "z" : Tensor[B x D x T]
-                Quantized continuous representation of input
-            "codes" : Tensor[B x N x T]
-                Codebook indices for each codebook
-                (quantized discrete representation of input)
-            "latents" : Tensor[B x N*D x T]
-                Projected latents (continuous representation of input before quantization)
-            "vq/commitment_loss" : Tensor[1]
-                Commitment loss to train encoder to predict vectors closer to codebook
-                entries
-            "vq/codebook_loss" : Tensor[1]
-                Codebook loss to update the codebook
-            "length" : int
-                Number of samples in input audio
-            "audio" : Tensor[B x 1 x length]
-                Decoded audio data.
-        """
-        length = audio_data.shape[-1]
-        audio_data = self.preprocess(audio_data, sample_rate)
-        vq_results = self.encode(audio_data, n_quantizers, **kwargs)
-        z = vq_results[0] if isinstance(vq_results, tuple) else vq_results.z
-        x = self.decode(z)
-        return x[..., :length], vq_results

fish_speech/models/dac/rvq.py

@@ -1,371 +0,0 @@
-import math
-import typing as tp
-from dataclasses import dataclass
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from fish_speech.models.dac.dac_nn import ResidualVectorQuantize
-from torch.nn.utils.parametrizations import weight_norm
-from torch.nn.utils.parametrize import remove_parametrizations
-
-
-def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]):
-    """Remove padding from x, handling properly zero padding. Only for 1d!"""
-    padding_left, padding_right = paddings
-    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
-    assert (padding_left + padding_right) <= x.shape[-1]
-    end = x.shape[-1] - padding_right
-    return x[..., padding_left:end]
-
-
-def get_extra_padding_for_conv1d(
-    x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0
-) -> int:
-    """See `pad_for_conv1d`."""
-    length = x.shape[-1]
-    n_frames = (length - kernel_size + padding_total) / stride + 1
-    ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
-    return ideal_length - length
-
-
-def pad1d(
-    x: torch.Tensor,
-    paddings: tp.Tuple[int, int],
-    mode: str = "zeros",
-    value: float = 0.0,
-):
-    """Tiny wrapper around F.pad, just to allow for reflect padding on small input.
-    If this is the case, we insert extra 0 padding to the right
-    before the reflection happen.
-    """
-    length = x.shape[-1]
-    padding_left, padding_right = paddings
-    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
-    if mode == "reflect":
-        max_pad = max(padding_left, padding_right)
-        extra_pad = 0
-        if length <= max_pad:
-            extra_pad = max_pad - length + 1
-            x = F.pad(x, (0, extra_pad))
-        padded = F.pad(x, paddings, mode, value)
-        end = padded.shape[-1] - extra_pad
-        return padded[..., :end]
-    else:
-        return F.pad(x, paddings, mode, value)
-
-
-class CausalConvNet(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        kernel_size,
-        dilation=1,
-        stride=1,
-        groups=1,
-        padding=None,
-    ):
-        super(CausalConvNet, self).__init__()
-        self.conv = nn.Conv1d(
-            in_channels,
-            out_channels,
-            kernel_size,
-            stride=stride,
-            dilation=dilation,
-            groups=groups,
-        )
-        self.stride = stride
-        self.kernel_size = (kernel_size - 1) * dilation + 1
-        self.dilation = dilation
-        self.padding = self.kernel_size - self.stride
-
-    def forward(self, x):
-        pad = self.padding
-        extra_padding = get_extra_padding_for_conv1d(
-            x, self.kernel_size, self.stride, pad
-        )
-        x = pad1d(x, (pad, extra_padding), mode="constant", value=0)
-        return self.conv(x).contiguous()
-
-    def weight_norm(self, name="weight", dim=0):
-        self.conv = weight_norm(self.conv, name=name, dim=dim)
-        return self
-
-    def remove_weight_norm(self):
-        self.conv = remove_parametrizations(self.conv)
-        return self
-
-
-class CausalTransConvNet(nn.Module):
-    def __init__(
-        self, in_channels, out_channels, kernel_size, dilation=1, stride=1, padding=None
-    ):
-        super(CausalTransConvNet, self).__init__()
-        self.conv = nn.ConvTranspose1d(
-            in_channels, out_channels, kernel_size, stride=stride, dilation=dilation
-        )
-        self.stride = stride
-        self.kernel_size = kernel_size
-
-    def forward(self, x):
-        x = self.conv(x)
-        pad = self.kernel_size - self.stride
-        padding_right = math.ceil(pad)
-        padding_left = pad - padding_right
-        x = unpad1d(x, (padding_left, padding_right))
-        return x.contiguous()
-
-    def weight_norm(self, name="weight", dim=0):
-        self.conv = weight_norm(self.conv, name=name, dim=dim)
-        return self
-
-    def remove_weight_norm(self):
-        self.conv = remove_parametrizations(self.conv)
-        return self
-
-
-# ConvNeXt Block copied from https://github.com/fishaudio/fish-diffusion/blob/main/fish_diffusion/modules/convnext.py
-class ConvNeXtBlock(nn.Module):
-    r"""ConvNeXt Block. There are two equivalent implementations:
-    (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
-    (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
-    We use (2) as we find it slightly faster in PyTorch
-    Args:
-        dim (int): Number of input channels.
-        drop_path (float): Stochastic depth rate. Default: 0.0
-        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.0.
-        kernel_size (int): Kernel size for depthwise conv. Default: 7.
-        dilation (int): Dilation for depthwise conv. Default: 1.
-    """  # noqa: E501
-
-    def __init__(
-        self,
-        dim: int,
-        layer_scale_init_value: float = 1e-6,
-        mlp_ratio: float = 4.0,
-        kernel_size: int = 7,
-        dilation: int = 1,
-    ):
-        super().__init__()
-        convnet_type = CausalConvNet
-        self.dwconv = convnet_type(
-            dim,
-            dim,
-            kernel_size=kernel_size,
-            # padding=int(dilation * (kernel_size - 1) / 2),
-            groups=dim,
-            dilation=dilation,
-        )  # depthwise conv
-        self.norm = nn.LayerNorm(dim, eps=1e-6)
-        self.pwconv1 = nn.Linear(
-            dim, int(mlp_ratio * dim)
-        )  # pointwise/1x1 convs, implemented with linear layers
-        self.act = nn.GELU()
-        self.pwconv2 = nn.Linear(int(mlp_ratio * dim), dim)
-        self.gamma = (
-            nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
-            if layer_scale_init_value > 0
-            else None
-        )
-
-    def forward(self, x, apply_residual: bool = True):
-        input = x
-
-        x = self.dwconv(x)
-        x = x.permute(0, 2, 1)  # (N, C, L) -> (N, L, C)
-        x = self.norm(x)
-        x = self.pwconv1(x)
-        x = self.act(x)
-        x = self.pwconv2(x)
-
-        if self.gamma is not None:
-            x = self.gamma * x
-
-        x = x.permute(0, 2, 1)  # (N, L, C) -> (N, C, L)
-
-        if apply_residual:
-            x = input + x
-
-        return x
-
-
-@dataclass
-class VQResult:
-    z: torch.Tensor
-    codes: torch.Tensor
-    latents: torch.Tensor
-    codebook_loss: torch.Tensor
-    commitment_loss: torch.Tensor
-    semantic_distill_z: torch.Tensor | None = None
-
-
-class DownsampleResidualVectorQuantize(nn.Module):
-    def __init__(
-        self,
-        input_dim: int = 1024,
-        n_codebooks: int = 9,
-        codebook_dim: int = 8,
-        quantizer_dropout: float = 0.5,
-        codebook_size: int = 1024,
-        semantic_codebook_size: int = 4096,
-        downsample_factor: tuple[int] = (2, 2),
-        downsample_dims: tuple[int] | None = None,
-        pre_module: nn.Module | None = None,
-        post_module: nn.Module | None = None,
-        semantic_predictor_module: nn.Module | None = None,
-    ):
-        super().__init__()
-
-        if downsample_dims is None:
-            downsample_dims = [input_dim for _ in range(len(downsample_factor))]
-
-        all_dims = (input_dim,) + tuple(downsample_dims)
-
-        self.semantic_quantizer = ResidualVectorQuantize(
-            input_dim=input_dim,
-            n_codebooks=1,
-            codebook_size=semantic_codebook_size,
-            codebook_dim=codebook_dim,
-            quantizer_dropout=0.0,
-        )
-
-        self.quantizer = ResidualVectorQuantize(
-            input_dim=input_dim,
-            n_codebooks=n_codebooks,
-            codebook_size=codebook_size,
-            codebook_dim=codebook_dim,
-            quantizer_dropout=quantizer_dropout,
-        )
-
-        self.downsample_factor = downsample_factor
-        self.downsample_dims = downsample_dims
-
-        convnet_type = CausalConvNet
-        transconvnet_type = CausalTransConvNet
-
-        self.downsample = nn.Sequential(
-            *[
-                nn.Sequential(
-                    convnet_type(
-                        all_dims[idx],
-                        all_dims[idx + 1],
-                        kernel_size=factor,
-                        stride=factor,
-                    ),
-                    ConvNeXtBlock(dim=all_dims[idx + 1]),
-                )
-                for idx, factor in enumerate(downsample_factor)
-            ]
-        )
-
-        self.upsample = nn.Sequential(
-            *[
-                nn.Sequential(
-                    transconvnet_type(
-                        all_dims[idx + 1],
-                        all_dims[idx],
-                        kernel_size=factor,
-                        stride=factor,
-                    ),
-                    ConvNeXtBlock(dim=all_dims[idx]),
-                )
-                for idx, factor in reversed(list(enumerate(downsample_factor)))
-            ]
-        )
-        self.apply(self._init_weights)
-        self.pre_module = (
-            pre_module if pre_module is not None else nn.Identity()
-        )  # leave for transformer, LSTM or Mamba or something else
-        self.post_module = post_module if post_module is not None else nn.Identity()
-        self.semantic_predictor_module = (
-            semantic_predictor_module
-            if semantic_predictor_module is not None
-            else nn.Identity()
-        )
-
-    def _init_weights(self, m):
-        if isinstance(m, (nn.Conv1d, nn.Linear)):
-            nn.init.trunc_normal_(m.weight, std=0.02)
-            nn.init.constant_(m.bias, 0)
-
-    def forward(
-        self, z, n_quantizers: int = None, semantic_len: torch.Tensor = None, **kwargs
-    ):
-        # z: (B, D, T)
-        original_shape = z.shape
-        if semantic_len is None:
-            semantic_len = torch.LongTensor([z.shape[-1]])
-        z = self.downsample(z)
-        z = self.pre_module(z)  # B, T, D
-        (
-            semantic_z,
-            semantic_codes,
-            semantic_latents,
-            semantic_commitment_loss,
-            semantic_codebook_loss,
-        ) = self.semantic_quantizer(z)
-        residual_z = z - semantic_z
-        residual_z, codes, latents, commitment_loss, codebook_loss = self.quantizer(
-            residual_z, n_quantizers=n_quantizers
-        )
-        z = semantic_z + residual_z
-        commitment_loss = commitment_loss + semantic_commitment_loss
-        codebook_loss = codebook_loss + semantic_codebook_loss
-        codes = torch.cat([semantic_codes, codes], dim=1)
-        latents = torch.cat([semantic_latents, latents], dim=1)
-        z = self.post_module(z)
-        z = self.upsample(z)
-        # z: (B, D, T)
-
-        # semantic distillation (disabled here since only used in training)
-        # semantic_distill_z = self.semantic_predictor_module(semantic_z, semantic_len).mT  # wav2vec target is B, T, D
-
-        # Pad or crop z to match original shape
-        diff = original_shape[-1] - z.shape[-1]
-        right = 0
-        left = abs(diff) - right
-
-        if diff > 0:
-            z = F.pad(z, (left, right))
-        elif diff < 0:
-            z = z[..., left:]
-
-        results = VQResult(
-            z=z,
-            codes=codes,
-            latents=latents,
-            commitment_loss=commitment_loss,
-            codebook_loss=codebook_loss,
-        )
-
-        return results
-
-    # def encode(self, z):
-    #     z = self.downsample(z)
-    #     z = self.pre_module(z)
-    #     _, indices, _, _, _ = self.quantizer(z.mT)
-    #     indices = rearrange(indices, "g b l r -> b (g r) l")
-    #     return indices
-    #
-    def decode(self, indices: torch.Tensor):
-        # indices = rearrange(indices, "b (g r) l -> g b l r", g=self.residual_fsq.groups)
-        indices[:, 0] = torch.clamp(
-            indices[:, 0], max=self.semantic_quantizer.codebook_size - 1
-        )
-        indices[:, 1:] = torch.clamp(
-            indices[:, 1:], max=self.quantizer.codebook_size - 1
-        )
-
-        z_q_semantic = self.semantic_quantizer.from_codes(indices[:, :1])[0]
-        z_q_residual = self.quantizer.from_codes(indices[:, 1:])[0]
-        z_q = z_q_semantic + z_q_residual
-        z_q = self.post_module(z_q)
-        z_q = self.upsample(z_q)
-        return z_q
-
-    # def from_latents(self, latents: torch.Tensor):
-    #     z_q, z_p, codes = super().from_latents(latents)
-    #     z_q = self.upsample(z_q)
-    #     return z_q, z_p, codes

fish_speech/models/text2semantic/__init__.py

@@ -1 +0,0 @@
-

fish_speech/models/text2semantic/errors.py

@@ -1,5 +0,0 @@
-class InvalidInputError(ValueError):
-    """The request cannot be served as given: bad text, references or length.
-
-    The API server answers these with 400; every other failure is a 500.
-    """

fish_speech/models/text2semantic/fish_omni/__init__.py

@@ -1 +0,0 @@
-

fish_speech/models/text2semantic/fish_omni/backbone.py

@@ -1,397 +0,0 @@
-import re
-from collections import defaultdict
-from collections.abc import Iterable, Sequence
-
-import torch
-import torch.nn as nn
-from loguru import logger
-from torch import Tensor
-from torch.types import Device
-
-from fish_speech.models.text2semantic.fish_omni.config import FishOmniTextConfig
-from fish_speech.models.text2semantic.fish_omni.fp8_moe import FP8_DTYPE
-from fish_speech.models.text2semantic.fish_omni.layers import (
-    FullAttention,
-    GatedDeltaNet,
-    LayerContext,
-    ZeroCenteredRMSNorm,
-)
-from fish_speech.models.text2semantic.fish_omni.moe import (
-    SparseMoeBlock,
-    clamp_to_float16,
-    dequantize_fp8_blocks,
-)
-
-CHECKPOINT_BACKBONE_PREFIX = "model.language_model."
-CHECKPOINT_LM_HEAD_KEY = "lm_head.weight"
-# Per-expert checkpoint tensors, named after the backbone prefix became "model.".
-EXPERT_KEY_PATTERN = re.compile(
-    r"model\.layers\.(?P<layer>\d+)\.mlp\.experts\.(?P<expert>\d+)\."
-    r"(?P<proj>gate_proj|up_proj|down_proj)\.(?P<kind>weight|weight_scale_inv)"
-)
-# Hugging Face's fused expert tensors (unquantized checkpoints).
-FUSED_EXPERT_KEY_PATTERN = re.compile(
-    r"model\.layers\.(?P<layer>\d+)\.mlp\.experts\.(?P<proj>gate_up_proj|down_proj)"
-)
-STACKED_EXPERT_PARAMETER_PATTERN = re.compile(
-    r"(model\.layers\.\d+\.mlp\.experts)\.(w13|w2)_(weight|scale_inv)"
-)
-
-
-class DecoderLayer(nn.Module):
-    def __init__(
-        self,
-        config: FishOmniTextConfig,
-        layer_index: int,
-        *,
-        fp8: bool,
-        dtype: torch.dtype | None = None,
-        device: Device = None,
-    ) -> None:
-        super().__init__()
-        self.layer_type = config.layer_types[layer_index]
-        if self.layer_type == "full_attention":
-            self.self_attn = FullAttention(config, dtype=dtype, device=device)
-        else:
-            self.linear_attn = GatedDeltaNet(config, dtype=dtype, device=device)
-        self.mlp = SparseMoeBlock(config, fp8=fp8, dtype=dtype, device=device)
-        norm = dict(dtype=dtype, device=device)
-        self.input_layernorm = ZeroCenteredRMSNorm(
-            config.hidden_size, config.rms_norm_eps, **norm
-        )
-        self.post_attention_layernorm = ZeroCenteredRMSNorm(
-            config.hidden_size, config.rms_norm_eps, **norm
-        )
-
-    def forward(self, x: Tensor, context: LayerContext) -> Tensor:
-        """x is the residual stream, possibly wider than the compute dtype."""
-        compute_dtype = self.input_layernorm.weight.dtype
-        mixer = (
-            self.self_attn if self.layer_type == "full_attention" else self.linear_attn
-        )
-        x = x + mixer(self.input_layernorm(x).to(compute_dtype), context)
-        return x + self.mlp(self.post_attention_layernorm(x).to(compute_dtype))
-
-
-class FishOmniTextModel(nn.Module):
-    def __init__(
-        self,
-        config: FishOmniTextConfig,
-        *,
-        fp8: bool,
-        dtype: torch.dtype | None = None,
-        device: Device = None,
-    ) -> None:
-        super().__init__()
-        self.embed_tokens = nn.Embedding(
-            config.vocab_size, config.hidden_size, dtype=dtype, device=device
-        )
-        self.layers = nn.ModuleList(
-            DecoderLayer(config, index, fp8=fp8, dtype=dtype, device=device)
-            for index in range(config.num_hidden_layers)
-        )
-        self.norm = ZeroCenteredRMSNorm(
-            config.hidden_size, config.rms_norm_eps, dtype=dtype, device=device
-        )
-
-
-class FishOmniBackbone(nn.Module):
-    """Qwen3.5-MoE hybrid text backbone + lm_head; layers may span several devices."""
-
-    def __init__(
-        self,
-        config: FishOmniTextConfig,
-        *,
-        fp8_experts: bool = False,
-        dtype: torch.dtype = torch.bfloat16,
-        device: Device = None,
-    ) -> None:
-        super().__init__()
-        self.config = config
-        self.dtype = dtype
-        # The residual stream of this model exceeds the float16 range.
-        self.residual_dtype = torch.float32 if dtype == torch.float16 else dtype
-        self.model = FishOmniTextModel(
-            config, fp8=fp8_experts, dtype=dtype, device=device
-        )
-        self.lm_head = nn.Linear(
-            config.hidden_size,
-            config.vocab_size,
-            bias=False,
-            dtype=dtype,
-            device=device,
-        )
-        self.max_batch_size = 0
-        self.max_seq_len = 0
-        self.inverse_frequencies: dict[torch.device, Tensor] = {}
-
-    def materialize(self, devices: Sequence[Device]) -> None:
-        """Allocate a meta-device backbone over devices in contiguous layer ranges.
-
-        Ranges are balanced by bytes, counting the embedding on the first device and
-        the final norm + lm_head on the last.
-        """
-        devices = [torch.device(d) for d in devices]
-        layer_bytes = [module_bytes(layer) for layer in self.model.layers]
-        head = module_bytes(self.model.embed_tokens)
-        tail = module_bytes(self.lm_head) + module_bytes(self.model.norm)
-        per_device = (head + sum(layer_bytes) + tail) / len(devices)
-        self.model.embed_tokens.to_empty(device=devices[0])
-        self.model.norm.to_empty(device=devices[-1])
-        self.lm_head.to_empty(device=devices[-1])
-        placement = defaultdict(list)
-        cumulative = head
-        for index, (layer, size) in enumerate(zip(self.model.layers, layer_bytes)):
-            middle = cumulative + size / 2
-            device = devices[min(int(middle // per_device), len(devices) - 1)]
-            layer.to_empty(device=device)
-            placement[str(device)].append(index)
-            cumulative += size
-        self.inverse_frequencies.clear()
-        logger.info(
-            "FishOmni backbone placement: "
-            + ", ".join(f"{d}: layers {v[0]}-{v[-1]}" for d, v in placement.items())
-            + f"; embed on {devices[0]}, lm_head on {devices[-1]}"
-        )
-
-    def setup_caches(self, max_batch_size: int, max_seq_len: int) -> None:
-        assert max_seq_len <= self.config.max_position_embeddings
-        for layer in self.model.layers:
-            device = layer.input_layernorm.weight.device
-            if layer.layer_type == "full_attention":
-                layer.self_attn.allocate_cache(
-                    max_batch_size, max_seq_len, self.dtype, device
-                )
-            else:
-                layer.linear_attn.allocate_cache(max_batch_size, self.dtype, device)
-        self.max_batch_size = max_batch_size
-        self.max_seq_len = max_seq_len
-
-    def embed(self, input_ids: Tensor) -> Tensor:
-        embeddings = self.model.embed_tokens
-        return embeddings(input_ids.to(embeddings.weight.device))
-
-    def forward(
-        self,
-        input_ids: Tensor | None = None,
-        *,
-        inputs_embeds: Tensor | None = None,
-        input_pos: Tensor | None = None,
-        static_kv_len: int | None = None,
-    ) -> Tensor:
-        """Final-norm hidden states [B, T, H] of B rows at the T positions input_pos.
-
-        With caches, rows at positions > 0 continue the cached sequence; without,
-        only a whole sequence from position 0 can run. static_kv_len (CUDA-graph
-        decode) attends that many cache rows without a host sync.
-        """
-        assert (input_ids is None) != (inputs_embeds is None)
-        hidden = self.embed(input_ids) if inputs_embeds is None else inputs_embeds
-        hidden = hidden.to(self.residual_dtype)
-        batch, length = hidden.shape[:2]
-        if static_kv_len is not None:
-            assert input_pos is not None and length == 1
-            assert 0 < static_kv_len <= self.max_seq_len
-            return self.run_layers(
-                hidden,
-                input_pos,
-                kv_len=static_kv_len,
-                continue_state=True,
-                use_cache=True,
-            )
-        if input_pos is None:
-            input_pos = torch.arange(length, device=hidden.device)
-        assert input_pos.shape == (length,)
-        start, end = int(input_pos[0]), int(input_pos[-1]) + 1
-        assert start >= 0 and end - start == length, "input_pos must be contiguous"
-        use_cache = self.max_seq_len > 0
-        if use_cache:
-            assert (
-                batch <= self.max_batch_size and end <= self.max_seq_len
-            ), "rows exceed the caches"
-        else:
-            assert start == 0, "continuing a sequence needs setup_caches first"
-        return self.run_layers(
-            hidden, input_pos, kv_len=end, continue_state=start > 0, use_cache=use_cache
-        )
-
-    def run_layers(
-        self,
-        hidden: Tensor,
-        input_pos: Tensor,
-        *,
-        kv_len: int,
-        continue_state: bool,
-        use_cache: bool,
-    ) -> Tensor:
-        rotary_dim = self.config.rotary_dim
-        contexts: dict[torch.device, LayerContext] = {}
-        for layer in self.model.layers:
-            device = layer.input_layernorm.weight.device
-            if device not in contexts:
-                frequencies = self.inverse_frequencies.get(device)
-                if frequencies is None:
-                    exponents = torch.arange(
-                        0, rotary_dim, 2, dtype=torch.float32, device=device
-                    )
-                    frequencies = 1.0 / (
-                        self.config.rope_theta ** (exponents / rotary_dim)
-                    )
-                    self.inverse_frequencies[device] = frequencies
-                positions = input_pos.to(device)
-                angles = positions.float()[:, None] * frequencies[None, :]
-                angles = torch.cat((angles, angles), dim=-1)
-                contexts[device] = LayerContext(
-                    input_pos=positions,
-                    cos=angles.cos(),
-                    sin=angles.sin(),
-                    kv_len=kv_len,
-                    continue_state=continue_state,
-                    use_cache=use_cache,
-                )
-            hidden = layer(hidden.to(device), contexts[device])
-        norm = self.model.norm
-        return norm(hidden.to(norm.weight.device)).to(self.dtype)
-
-    def compute_logits(self, hidden: Tensor) -> Tensor:
-        return self.lm_head(hidden.to(self.lm_head.weight.device))
-
-
-def module_bytes(module: nn.Module) -> int:
-    tensors = [*module.parameters(), *module.buffers()]
-    return sum(t.numel() * t.element_size() for t in tensors)
-
-
-def backbone_name(checkpoint_name: str) -> str:
-    """Map a checkpoint key onto a FishOmniBackbone parameter name."""
-    if checkpoint_name == CHECKPOINT_LM_HEAD_KEY:
-        return "lm_head.weight"
-    if checkpoint_name.startswith(CHECKPOINT_BACKBONE_PREFIX):
-        return "model." + checkpoint_name[len(CHECKPOINT_BACKBONE_PREFIX) :]
-    if checkpoint_name.startswith("model."):
-        return checkpoint_name
-    raise ValueError(f"checkpoint key {checkpoint_name!r} is not a backbone weight")
-
-
-@torch.no_grad()
-def load_backbone_weights(
-    model: FishOmniBackbone, weights: Iterable[tuple[str, Tensor]]
-) -> None:
-    """Copy checkpoint tensors into model, filling every parameter slot exactly once.
-
-    Slots are (name, expert, part): stacked expert tensors are filled per expert and
-    per gate / up / down projection, everything else whole.
-    """
-    targets = dict(model.named_parameters())
-    targets.update(model.named_buffers())
-    expected: set[tuple[str, int, str]] = set()
-    for name in targets:
-        match = STACKED_EXPERT_PARAMETER_PATTERN.fullmatch(name)
-        if match is None:
-            expected.add((name, -1, "whole"))
-            continue
-        experts = model.get_submodule(match.group(1))
-        parts = ("gate", "up") if match.group(2) == "w13" else ("down",)
-        for expert in range(experts.num_experts):
-            expected.update((name, expert, part) for part in parts)
-    loaded: set[tuple[str, int, str]] = set()
-    # FP8 checkpoint experts loaded into dense experts: a weight or scale waits here
-    # for its partner, then the pair is dequantized on the target's device.
-    pending_fp8: dict[tuple[str, int, str], dict[str, Tensor]] = {}
-
-    def fill(name: str, expert: int, part: str, target: Tensor, tensor: Tensor) -> None:
-        slot = (name, expert, part)
-        if slot not in expected:
-            raise ValueError(f"no parameter slot {slot}")
-        if slot in loaded:
-            raise ValueError(f"parameter slot {slot} filled twice")
-        if target.shape != tensor.shape:
-            raise ValueError(
-                f"{slot}: checkpoint shape {tuple(tensor.shape)} != "
-                f"{tuple(target.shape)}"
-            )
-        if target.dtype == FP8_DTYPE and tensor.dtype != FP8_DTYPE:
-            raise ValueError(f"{slot}: expected {FP8_DTYPE}, got {tensor.dtype}")
-        if target.dtype == torch.float16 and tensor.dtype != torch.float16:
-            tensor = clamp_to_float16(tensor, f"{name}[{expert}].{part}")
-        target.copy_(tensor)
-        loaded.add(slot)
-
-    for checkpoint_name, tensor in weights:
-        name = backbone_name(checkpoint_name)
-        expert_match = EXPERT_KEY_PATTERN.fullmatch(name)
-        fused_match = FUSED_EXPERT_KEY_PATTERN.fullmatch(name)
-        if expert_match is not None:
-            prefix = f"model.layers.{expert_match['layer']}.mlp.experts"
-            experts = model.get_submodule(prefix)
-            expert = int(expert_match["expert"])
-            projection = expert_match["proj"]
-            is_scale = expert_match["kind"] == "weight_scale_inv"
-            if not is_scale and experts.fp8 and tensor.dtype != FP8_DTYPE:
-                raise ValueError(f"{checkpoint_name}: expected an FP8 expert weight")
-            if not experts.fp8 and (is_scale or tensor.dtype == FP8_DTYPE):
-                key = (prefix, expert, projection)
-                halves = pending_fp8.setdefault(key, {})
-                if expert_match["kind"] in halves:
-                    raise ValueError(f"{checkpoint_name} appears twice")
-                halves[expert_match["kind"]] = tensor
-                if len(halves) < 2:
-                    continue
-                del pending_fp8[key]
-                target = experts.w13_weight
-                tensor = dequantize_fp8_blocks(
-                    halves["weight"].to(target.device),
-                    halves["weight_scale_inv"].to(target.device),
-                    torch.float32,
-                )
-                if target.dtype == torch.float16:
-                    tensor = clamp_to_float16(tensor, checkpoint_name)
-                tensor = tensor.to(target.dtype)
-                is_scale = False
-            suffix = "scale_inv" if is_scale else "weight"
-            if projection == "down_proj":
-                full_name = f"{prefix}.w2_{suffix}"
-                fill(full_name, expert, "down", targets[full_name][expert], tensor)
-            else:
-                full_name = f"{prefix}.w13_{suffix}"
-                rows = targets[full_name].shape[1] // 2
-                gate = projection == "gate_proj"
-                half = slice(0, rows) if gate else slice(rows, 2 * rows)
-                fill(
-                    full_name,
-                    expert,
-                    "gate" if gate else "up",
-                    targets[full_name][expert, half],
-                    tensor,
-                )
-        elif fused_match is not None:
-            prefix = f"model.layers.{fused_match['layer']}.mlp.experts"
-            if fused_match["proj"] == "down_proj":
-                full_name = f"{prefix}.w2_weight"
-                for expert, matrix in enumerate(tensor):
-                    fill(full_name, expert, "down", targets[full_name][expert], matrix)
-            else:
-                full_name = f"{prefix}.w13_weight"
-                rows = targets[full_name].shape[1] // 2
-                for expert, matrix in enumerate(tensor):
-                    target = targets[full_name][expert]
-                    fill(full_name, expert, "gate", target[:rows], matrix[:rows])
-                    fill(full_name, expert, "up", target[rows:], matrix[rows:])
-        else:
-            if name not in targets:
-                raise ValueError(
-                    f"checkpoint key {checkpoint_name!r} has no parameter {name!r}"
-                )
-            fill(name, -1, "whole", targets[name], tensor)
-
-    if pending_fp8:
-        raise ValueError(
-            f"{len(pending_fp8)} FP8 experts lack a weight or scale, "
-            f"e.g. {sorted(pending_fp8)[:4]}"
-        )
-    if missing := expected - loaded:
-        raise ValueError(
-            f"{len(missing)} backbone parameter slots were not loaded, "
-            f"e.g. {sorted(missing)[:8]}"
-        )
-    logger.info(f"FishOmni backbone: loaded {len(loaded)} parameter slots")

fish_speech/models/text2semantic/fish_omni/chunking.py

@@ -1,206 +0,0 @@
-"""Long-text chunking for Drama 3, modeled on the hosted inference server.
-
-Text is cut into chunks of about DRAMA3_CHUNK_BYTES UTF-8 bytes with at most
-DRAMA3_CHUNK_MAX_SPEAKERS speakers, preferring sentence ends, then clause marks, then
-spaces. Each chunk is generated on its own; the first chunk and the latest one go
-back into the next prompt as user text plus assistant audio, which keeps the voice
-consistent across chunks.
-"""
-
-import re
-from collections import deque
-from collections.abc import Iterator
-
-import torch
-
-from fish_speech.models.text2semantic.fish_omni.generate import (
-    SPEAKER_TAG_PATTERN,
-    UNICODE_WHITE_SPACE,
-)
-
-# A threshold, not a hard cap: a protected span longer than this stays whole.
-DRAMA3_CHUNK_BYTES = 300
-DRAMA3_CHUNK_MAX_SPEAKERS = 3
-# Frames per chunk for max_new_tokens <= 0, about 27.9 s at 2048 samples per
-# frame and 44.1 kHz.
-DRAMA3_CHUNK_MAX_NEW_TOKENS = 600
-# History besides the pinned first chunk.
-DRAMA3_HISTORY_RECENT_TURNS = 1
-
-SENTENCE_ENDS = frozenset(".!?;。！？；…\n")
-CLAUSE_MARKS = frozenset(",:，：、")
-SPACES = frozenset(" ")
-# Never split inside: special tokens, paired tags such as <whisper>…</whisper>, and
-# [bracketed] cues.
-PROTECTED_SPAN = re.compile(
-    r"<\|[^|>]*\|>|<([A-Za-z][\w-]*)\b[^>]*>.*?</\1>|\[[^\[\]]*\]", re.S
-)
-
-Piece = tuple[str, bool]  # (text, protected)
-HistoryTurn = tuple[str, torch.Tensor]  # (chunk text, codes [num_codebooks, frames])
-
-
-def utf8_len(text: str) -> int:
-    return len(text.encode("utf-8"))
-
-
-def _protected_pieces(text: str) -> list[Piece]:
-    pieces, last = [], 0
-    for match in PROTECTED_SPAN.finditer(text):
-        if match.start() > last:
-            pieces.append((text[last : match.start()], False))
-        pieces.append((match.group(0), True))
-        last = match.end()
-    if last < len(text):
-        pieces.append((text[last:], False))
-    return pieces
-
-
-def _breaks_at(text: str, index: int, marks: frozenset[str]) -> bool:
-    if text[index] not in marks:
-        return False
-    # Keep decimals such as 3.14 in one piece.
-    return not (
-        text[index] == "."
-        and 0 < index < len(text) - 1
-        and text[index - 1].isdigit()
-        and text[index + 1].isdigit()
-    )
-
-
-def _break_at_marks(pieces: list[Piece], length: int, marks: frozenset[str]):
-    """Split free pieces longer than length bytes after each mark, keeping it."""
-    out = []
-    for text, protected in pieces:
-        if protected or utf8_len(text) <= length:
-            out.append((text, protected))
-            continue
-        start = 0
-        for index in range(len(text)):
-            if _breaks_at(text, index, marks):
-                out.append((text[start : index + 1], False))
-                start = index + 1
-        if start < len(text):
-            out.append((text[start:], False))
-    return out
-
-
-def _break_by_length(pieces: list[Piece], length: int) -> list[Piece]:
-    """Split the free pieces still longer than length on character boundaries."""
-    out = []
-    for text, protected in pieces:
-        if protected or utf8_len(text) <= length:
-            out.append((text, protected))
-            continue
-        current, size = "", 0
-        for char in text:
-            if current and size + utf8_len(char) > length:
-                out.append((current, False))
-                current, size = "", 0
-            current += char
-            size += utf8_len(char)
-        out.append((current, False))
-    return out
-
-
-def split_by_length(text: str, length: int = DRAMA3_CHUNK_BYTES) -> list[str]:
-    """Chunks of about length bytes, packed sentence by sentence.
-
-    Text is cut at every sentence end; a sentence still longer than length is cut
-    at clause marks, then spaces, then on raw length. Adjacent pieces are then
-    merged back up to length.
-    """
-    if not text.strip(UNICODE_WHITE_SPACE):
-        return []
-    if utf8_len(text) <= length:
-        return [text]
-    pieces = _break_at_marks(_protected_pieces(text), 0, SENTENCE_ENDS)
-    for marks in (CLAUSE_MARKS, SPACES):
-        pieces = _break_at_marks(pieces, length, marks)
-    pieces = _break_by_length(pieces, length)
-
-    chunks, current = [], ""
-    for piece, _ in pieces:
-        if current and utf8_len(current + piece) > length:
-            chunks.append(current)
-            current = ""
-        current += piece
-    chunks.append(current)
-    return [
-        chunk.strip(UNICODE_WHITE_SPACE)
-        for chunk in chunks
-        if chunk.strip(UNICODE_WHITE_SPACE)
-    ]
-
-
-def parse_speaker_segments(text: str) -> list[tuple[str, str]]:
-    """(speaker id, text) per tagged segment; text before the first tag is speaker 0."""
-    parts = SPEAKER_TAG_PATTERN.split(text)
-    segments = [("0", parts[0])] if parts[0].strip(UNICODE_WHITE_SPACE) else []
-    for speaker, body in zip(parts[1::2], parts[2::2]):
-        if body.strip(UNICODE_WHITE_SPACE):
-            segments.append((speaker, body))
-    return segments
-
-
-def split_drama3_text(
-    text: str,
-    max_bytes: int = DRAMA3_CHUNK_BYTES,
-    max_speakers: int = DRAMA3_CHUNK_MAX_SPEAKERS,
-) -> list[str]:
-    """Generation chunks for text, in order.
-
-    Untagged text is split by length. Tagged text is split into speaker segments
-    (long ones cut further, each piece keeping its tag) and regrouped: a new chunk
-    starts when a segment would bring in one speaker too many or exceed max_bytes,
-    so A, B, A can share a chunk.
-    """
-    if SPEAKER_TAG_PATTERN.search(text) is None:
-        return split_by_length(text, max_bytes)
-
-    segments = []
-    for speaker, body in parse_speaker_segments(text):
-        for piece in split_by_length(body, max_bytes):
-            segments.append((speaker, "<" + f"|speaker:{speaker}|> {piece}"))
-
-    chunks, batch, speakers, size = [], [], [], 0
-    for speaker, segment in segments:
-        new_speaker = speaker not in speakers
-        too_many_speakers = new_speaker and len(speakers) >= max_speakers
-        too_long = bool(batch) and size + utf8_len(segment) > max_bytes
-        if too_many_speakers or too_long:
-            chunks.append("\n".join(batch))
-            batch, speakers, size = [], [], 0
-            new_speaker = True
-        batch.append(segment)
-        if new_speaker:
-            speakers.append(speaker)
-        size += utf8_len(segment)
-    if batch:
-        chunks.append("\n".join(batch))
-    return chunks
-
-
-class Drama3History:
-    """Context for the next chunk: the first chunk stays pinned, plus the latest
-    DRAMA3_HISTORY_RECENT_TURNS chunks."""
-
-    def __init__(self, recent_turns: int = DRAMA3_HISTORY_RECENT_TURNS) -> None:
-        self.pinned: HistoryTurn | None = None
-        self.recent: deque[HistoryTurn] = deque(maxlen=recent_turns)
-
-    def push(self, text: str, codes: torch.Tensor) -> None:
-        if self.pinned is None:
-            self.pinned = (text, codes)
-        else:
-            self.recent.append((text, codes))
-
-    def candidates(self) -> Iterator[list[HistoryTurn]]:
-        """Histories to try, longest first: drop the oldest recent turns, then the
-        pinned one, ending with no history."""
-        pinned = [self.pinned] if self.pinned is not None else []
-        recent = list(self.recent)
-        for start in range(len(recent) + 1):
-            yield pinned + recent[start:]
-        if pinned:
-            yield []

fish_speech/models/text2semantic/fish_omni/config.py

@@ -1,249 +0,0 @@
-from dataclasses import dataclass, fields
-from typing import Any
-
-from fish_speech.models.text2semantic.fish_omni.fp8_moe import FP8_BLOCK_SIZE
-
-FISH_OMNI_MODEL_TYPE = "fish_omni"
-TEXT_MODEL_TYPES = ("fish_omni_text", "qwen3_5_moe_text")
-SPEECH_TEMPLATE_VERSION = "reference-bracket-v1"
-FP8_QUANT_METHOD = "bsgemm-moe"
-# Text config flags fixed by the architecture: (key, default when absent, required value).
-FIXED_TEXT_FLAGS = (
-    ("rotary_interleaved", False, False),
-    ("use_qk_norm", True, True),
-    ("attn_output_gate", True, True),
-    ("layernorm_zero_centered_gamma", True, True),
-    ("shared_expert_gate", True, True),
-)
-# attention_qk_norm is deliberately not checked; see TowerAttention.
-FIXED_AUDIO_DECODER_FLAGS = {
-    "decoder_type": "codebook_mtp_shared",
-    "rotary_interleaved": True,
-    "attention_output_gate": True,
-    "layernorm_zero_centered_gamma": True,
-    "use_moe": False,
-    "mtp_use_swa": False,
-    "attention_qkv_bias": False,
-    "attention_o_bias": False,
-    "dropout": 0.0,
-}
-
-
-@dataclass(kw_only=True)
-class FishOmniTextConfig:
-    vocab_size: int
-    hidden_size: int
-    num_hidden_layers: int
-    layer_types: list[str]
-    num_attention_heads: int
-    num_key_value_heads: int
-    head_dim: int
-    rope_theta: float
-    partial_rotary_factor: float
-    rms_norm_eps: float
-    max_position_embeddings: int
-    linear_conv_kernel_dim: int
-    linear_key_head_dim: int
-    linear_value_head_dim: int
-    linear_num_key_heads: int
-    linear_num_value_heads: int
-    num_experts: int
-    num_experts_per_tok: int
-    moe_intermediate_size: int
-    shared_expert_intermediate_size: int
-    norm_topk_prob: bool = True
-
-    def __post_init__(self) -> None:
-        if len(self.layer_types) != self.num_hidden_layers or not set(
-            self.layer_types
-        ) <= {"linear_attention", "full_attention"}:
-            raise ValueError(
-                "layer_types must name linear_attention / full_attention for each of "
-                f"the {self.num_hidden_layers} layers; got {self.layer_types!r}"
-            )
-        if self.num_attention_heads % self.num_key_value_heads:
-            raise ValueError(
-                "num_attention_heads must be a multiple of num_key_value_heads"
-            )
-        if self.linear_num_value_heads % self.linear_num_key_heads:
-            raise ValueError(
-                "linear_num_value_heads must be a multiple of linear_num_key_heads"
-            )
-        if self.num_experts_per_tok > self.num_experts:
-            raise ValueError("num_experts_per_tok exceeds num_experts")
-        if self.rotary_dim <= 0 or self.rotary_dim % 2:
-            raise ValueError(
-                f"rotary width must be positive and even; got {self.rotary_dim}"
-            )
-
-    @property
-    def rotary_dim(self) -> int:
-        return int(self.head_dim * self.partial_rotary_factor)
-
-    @classmethod
-    def from_dict(cls, data: dict[str, Any]) -> "FishOmniTextConfig":
-        """Parse a fish_omni_text or qwen3_5_moe_text config."""
-        if data.get("model_type") not in TEXT_MODEL_TYPES:
-            raise ValueError(
-                f"text_config.model_type={data.get('model_type')!r}; "
-                f"expected one of {TEXT_MODEL_TYPES}"
-            )
-        if any(t != "moe" for t in data.get("mlp_layer_types") or ()):
-            raise NotImplementedError("every text_config.mlp_layer_types must be moe")
-        for key, default, required in FIXED_TEXT_FLAGS:
-            if data.get(key, default) is not required:
-                raise NotImplementedError(
-                    f"text_config.{key} must be {required}; got {data.get(key)!r}"
-                )
-        if data.get("tie_word_embeddings", False):
-            raise NotImplementedError("the backbone keeps a separate lm_head")
-        rope = data.get("rope_parameters") or data
-        if rope.get("rope_type", rope.get("type", "default")) != "default":
-            raise NotImplementedError("only the default rope_type is supported")
-
-        num_layers = data["num_hidden_layers"]
-        layer_types = data.get("layer_types") or [
-            (
-                "full_attention"
-                if (i + 1) % data["full_attention_interval"] == 0
-                else "linear_attention"
-            )
-            for i in range(num_layers)
-        ]
-        # Text-only positions are equal on every M-RoPE axis, so an mrope_section
-        # reduces to the plain 1-D rotation.
-        return cls(
-            vocab_size=data["vocab_size"],
-            hidden_size=data["hidden_size"],
-            num_hidden_layers=num_layers,
-            layer_types=list(layer_types),
-            num_attention_heads=data["num_attention_heads"],
-            num_key_value_heads=data["num_key_value_heads"],
-            head_dim=data["head_dim"],
-            rope_theta=rope["rope_theta"],
-            # Older configs keep the factor at the top level; 0.25 is the HF default.
-            partial_rotary_factor=rope.get(
-                "partial_rotary_factor", data.get("partial_rotary_factor", 0.25)
-            ),
-            rms_norm_eps=data["rms_norm_eps"],
-            max_position_embeddings=data["max_position_embeddings"],
-            linear_conv_kernel_dim=data["linear_conv_kernel_dim"],
-            linear_key_head_dim=data["linear_key_head_dim"],
-            linear_value_head_dim=data["linear_value_head_dim"],
-            linear_num_key_heads=data["linear_num_key_heads"],
-            linear_num_value_heads=data["linear_num_value_heads"],
-            num_experts=data["num_experts"],
-            num_experts_per_tok=data["num_experts_per_tok"],
-            moe_intermediate_size=data["moe_intermediate_size"],
-            shared_expert_intermediate_size=data["shared_expert_intermediate_size"],
-            norm_topk_prob=data.get("norm_topk_prob", True),
-        )
-
-
-@dataclass(kw_only=True, frozen=True)
-class AudioDecoderConfig:
-    dim: int
-    intermediate_size: int
-    n_head: int
-    n_local_heads: int
-    head_dim: int
-    norm_eps: float
-    rope_base: float
-    rotary_percent: float
-    vocab_size: int
-    num_codebooks: int
-    shared_enrichment_layers: int
-    vq_backbone_codebooks: int
-
-    def __post_init__(self) -> None:
-        if (
-            self.n_head % self.n_local_heads
-            or not 0 < self.rotary_dim <= self.head_dim
-            or self.rotary_dim % 2
-        ):
-            raise ValueError("invalid tower attention geometry")
-        if (
-            self.num_codebooks < 2
-            or not 0 <= self.vq_backbone_codebooks < self.num_codebooks
-        ):
-            raise ValueError("invalid codebook counts")
-        if self.shared_enrichment_layers < 1:
-            raise ValueError("shared_enrichment_layers must be >= 1")
-
-    @property
-    def rotary_dim(self) -> int:
-        return int(self.head_dim * self.rotary_percent)
-
-    @classmethod
-    def from_dict(cls, data: dict[str, Any]) -> "AudioDecoderConfig":
-        for key, value in FIXED_AUDIO_DECODER_FLAGS.items():
-            if data.get(key, value) != value:
-                raise ValueError(
-                    f"unsupported audio_decoder_config.{key}={data[key]!r}"
-                )
-        if data.get("text_dim", data["dim"]) != data["dim"]:
-            raise ValueError("audio_decoder_config.text_dim must equal dim")
-        if data.get("n_layer") != data["shared_enrichment_layers"]:
-            raise ValueError("n_layer must equal shared_enrichment_layers")
-        placement = data.get("vq_feature_placement")
-        if placement not in ("backbone_partial", "mtp_shared_pre") or (
-            placement == "mtp_shared_pre" and data["vq_backbone_codebooks"] != 0
-        ):
-            raise ValueError(f"unsupported vq_feature_placement {placement!r}")
-        return cls(**{f.name: data[f.name] for f in fields(cls)})
-
-
-@dataclass(kw_only=True)
-class FishOmniConfig:
-    text: FishOmniTextConfig
-    audio_decoder: AudioDecoderConfig
-    semantic_start_token_id: int
-    semantic_end_token_id: int
-    eos_token_id: int
-    fp8_experts: bool = False
-
-    def __post_init__(self) -> None:
-        start, end = self.semantic_start_token_id, self.semantic_end_token_id
-        if not 0 <= start <= end < self.text.vocab_size:
-            raise ValueError("the semantic band must lie inside the text vocabulary")
-        if start <= self.eos_token_id <= end:
-            raise ValueError("eos_token_id lies inside the semantic band")
-        if self.audio_decoder.dim != self.text.hidden_size:
-            raise ValueError("tower and backbone hidden sizes must match")
-        if self.audio_decoder.vocab_size != end - start + 1:
-            raise ValueError("tower vocabulary must match the semantic token band")
-
-    @classmethod
-    def from_dict(cls, data: dict[str, Any]) -> "FishOmniConfig":
-        if data.get("model_type") != FISH_OMNI_MODEL_TYPE:
-            raise ValueError(
-                f"model_type={data.get('model_type')!r}; "
-                f"expected {FISH_OMNI_MODEL_TYPE!r}"
-            )
-        if data.get("audio_encoder_config") is not None:
-            raise NotImplementedError("audio_encoder_config is not supported")
-        if data.get("speech_template_version") != SPEECH_TEMPLATE_VERSION:
-            raise NotImplementedError(
-                f"speech_template_version={data.get('speech_template_version')!r}; "
-                f"only {SPEECH_TEMPLATE_VERSION} is supported"
-            )
-        quantization = data.get("quantization_config")
-        if quantization is not None and (
-            quantization.get("quant_method") != FP8_QUANT_METHOD
-            or quantization.get("weight_block_size") != [FP8_BLOCK_SIZE] * 2
-        ):
-            raise NotImplementedError(
-                f"only {FP8_QUANT_METHOD} quantization with weight_block_size "
-                f"[{FP8_BLOCK_SIZE}, {FP8_BLOCK_SIZE}] is supported"
-            )
-        eos = data.get("eos_token_id", data["text_config"].get("eos_token_id"))
-        if eos is None:
-            raise ValueError("config has no eos_token_id")
-        return cls(
-            text=FishOmniTextConfig.from_dict(data["text_config"]),
-            audio_decoder=AudioDecoderConfig.from_dict(data["audio_decoder_config"]),
-            semantic_start_token_id=data["semantic_start_token_id"],
-            semantic_end_token_id=data["semantic_end_token_id"],
-            eos_token_id=eos,
-            fp8_experts=quantization is not None,
-        )

fish_speech/models/text2semantic/fish_omni/fp8_moe.py

@@ -1,355 +0,0 @@
-from collections.abc import Iterable
-
-import torch
-from torch import Tensor
-
-try:
-    import triton
-    import triton.language as tl
-except ImportError:  # Triton ships only with the CUDA / ROCm builds of torch.
-    triton = None
-
-FP8_DTYPE = torch.float8_e4m3fn
-# Weight block edge of bsgemm-moe checkpoints; also the activation quant group width.
-FP8_BLOCK_SIZE = 128
-FP8_MAX = float(torch.finfo(FP8_DTYPE).max)
-FP8_MIN_SCALE = 1e-10
-GEMM_BLOCK_N = 128
-GEMM_NUM_WARPS = 4
-GEMM_NUM_STAGES = 3
-
-
-def fp8_gemm_supported(devices: Iterable[str | torch.device]) -> bool:
-    """Every device is a CUDA SM89+ GPU (FP8 tensor cores) and Triton is available.
-
-    ROCm also reports its GPUs as "cuda" with a capability, but the kernels use the
-    e4m3fn format, which AMD GPUs lack (MI300 uses fnuz, RDNA3 has no FP8).
-    """
-    if triton is None or not torch.cuda.is_available() or torch.version.hip:
-        return False
-    devices = [torch.device(device) for device in devices]
-    return all(
-        device.type == "cuda" and torch.cuda.get_device_capability(device) >= (8, 9)
-        for device in devices
-    )
-
-
-if triton is not None:
-
-    @triton.jit
-    def quant_kernel(
-        x_ptr,
-        out_ptr,
-        scale_ptr,
-        stride_x,
-        cols,
-        fp8_max: tl.constexpr,
-        min_scale: tl.constexpr,
-        GROUP: tl.constexpr,
-    ):
-        """Dynamic FP8 quantization of one (row, channel group)."""
-        row = tl.program_id(0).to(tl.int64)
-        group = tl.program_id(1)
-        offs = group * GROUP + tl.arange(0, GROUP)
-        mask = offs < cols
-        x = tl.load(x_ptr + row * stride_x + offs, mask=mask, other=0.0)
-        x = x.to(tl.float32)
-        scale = tl.maximum(tl.max(tl.abs(x)), min_scale) / fp8_max
-        q = tl.clamp(x / scale, -fp8_max, fp8_max)
-        tl.store(out_ptr + row * cols + offs, q.to(out_ptr.dtype.element_ty), mask=mask)
-        tl.store(scale_ptr + row * tl.num_programs(1) + group, scale)
-
-    @triton.jit
-    def silu_mul_quant_kernel(
-        x_ptr,
-        out_ptr,
-        scale_ptr,
-        stride_x,
-        half,
-        fp8_max: tl.constexpr,
-        min_scale: tl.constexpr,
-        GROUP: tl.constexpr,
-    ):
-        """silu(gate) * up of x = [gate | up], then FP8 quantization per group."""
-        row = tl.program_id(0).to(tl.int64)
-        group = tl.program_id(1)
-        offs = group * GROUP + tl.arange(0, GROUP)
-        mask = offs < half
-        gate = tl.load(x_ptr + row * stride_x + offs, mask=mask, other=0.0)
-        up = tl.load(x_ptr + row * stride_x + half + offs, mask=mask, other=0.0)
-        gate, up = gate.to(tl.float32), up.to(tl.float32)
-        y = gate * tl.sigmoid(gate) * up
-        scale = tl.maximum(tl.max(tl.abs(y)), min_scale) / fp8_max
-        q = tl.clamp(y / scale, -fp8_max, fp8_max)
-        tl.store(out_ptr + row * half + offs, q.to(out_ptr.dtype.element_ty), mask=mask)
-        tl.store(scale_ptr + row * tl.num_programs(1) + group, scale)
-
-    @triton.jit
-    def fused_moe_kernel(
-        a_ptr,
-        b_ptr,
-        c_ptr,
-        a_scale_ptr,
-        b_scale_ptr,
-        topk_weights_ptr,
-        sorted_ids_ptr,
-        expert_ids_ptr,
-        num_post_padded_ptr,
-        N,
-        K,
-        num_valid,
-        stride_am,
-        stride_be,
-        stride_bn,
-        stride_bk,
-        stride_cm,
-        stride_asm,
-        stride_bse,
-        stride_bsn,
-        stride_bsk,
-        GROUP: tl.constexpr,
-        BLOCK_M: tl.constexpr,
-        BLOCK_N: tl.constexpr,
-        BLOCK_K: tl.constexpr,
-        GROUP_M: tl.constexpr,
-        TOP_K: tl.constexpr,
-        MUL_ROUTED_WEIGHT: tl.constexpr,
-    ):
-        """C[pair] = A[pair // TOP_K] @ B[expert(pair)]^T for one output tile.
-
-        Rows are (token, slot) pairs sorted by expert, BLOCK_M rows per expert;
-        padding rows hold num_valid and are masked out.
-        """
-        pid = tl.program_id(0)
-        EM = tl.load(num_post_padded_ptr)
-        num_pid_m = tl.cdiv(EM, BLOCK_M)
-        num_pid_n = tl.cdiv(N, BLOCK_N)
-        # The grid is sized for the worst case; blocks past the real layout exit.
-        if pid >= num_pid_m * num_pid_n:
-            return
-        num_pid_in_group = GROUP_M * num_pid_n
-        group_id = pid // num_pid_in_group
-        first_pid_m = group_id * GROUP_M
-        group_size_m = min(num_pid_m - first_pid_m, GROUP_M)
-        pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
-        pid_n = (pid % num_pid_in_group) // group_size_m
-
-        offs_row = pid_m * BLOCK_M + tl.arange(0, BLOCK_M).to(tl.int64)
-        pair = tl.load(sorted_ids_ptr + offs_row).to(tl.int64)
-        pair_mask = pair < num_valid
-        expert = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
-
-        offs_n = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N).to(tl.int64)) % N
-        offs_k = tl.arange(0, BLOCK_K)
-        token = pair // TOP_K
-        a_ptrs = a_ptr + token[:, None] * stride_am + offs_k[None, :]
-        b_ptrs = (
-            b_ptr
-            + expert * stride_be
-            + offs_k[:, None] * stride_bk
-            + offs_n[None, :] * stride_bn
-        )
-        a_scale_ptrs = a_scale_ptr + token * stride_asm
-        b_scale_ptrs = (
-            b_scale_ptr + expert * stride_bse + (offs_n // GROUP) * stride_bsn
-        )
-
-        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
-        for k in range(0, tl.cdiv(K, BLOCK_K)):
-            k_mask = offs_k < K - k * BLOCK_K
-            a = tl.load(a_ptrs, mask=pair_mask[:, None] & k_mask[None, :], other=0.0)
-            b = tl.load(b_ptrs, mask=k_mask[:, None], other=0.0)
-            k_group = (k * BLOCK_K) // GROUP
-            a_scale = tl.load(a_scale_ptrs + k_group, mask=pair_mask, other=0.0)
-            b_scale = tl.load(b_scale_ptrs + k_group * stride_bsk)
-            acc += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :]
-            a_ptrs += BLOCK_K
-            b_ptrs += BLOCK_K * stride_bk
-
-        if MUL_ROUTED_WEIGHT:
-            w = tl.load(topk_weights_ptr + pair, mask=pair_mask, other=0.0)
-            acc = acc * w.to(tl.float32)[:, None]
-        offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
-        c_ptrs = c_ptr + pair[:, None] * stride_cm + offs_cn[None, :]
-        tl.store(
-            c_ptrs,
-            acc.to(c_ptr.dtype.element_ty),
-            mask=pair_mask[:, None] & (offs_cn[None, :] < N),
-        )
-
-
-def quantize_groups(
-    kernel: "triton.JITFunction", x: Tensor, cols: int
-) -> tuple[Tensor, Tensor]:
-    """Run a per-token-group quant kernel: FP8 [M, cols] + fp32 scales [M, groups]."""
-    rows = x.shape[0]
-    groups = triton.cdiv(cols, FP8_BLOCK_SIZE)
-    quantized = torch.empty((rows, cols), dtype=FP8_DTYPE, device=x.device)
-    scales = torch.empty((rows, groups), dtype=torch.float32, device=x.device)
-    kernel[(rows, groups)](
-        x,
-        quantized,
-        scales,
-        x.stride(0),
-        cols,
-        fp8_max=FP8_MAX,
-        min_scale=FP8_MIN_SCALE,
-        GROUP=FP8_BLOCK_SIZE,
-        num_warps=1,
-    )
-    return quantized, scales
-
-
-def moe_align(
-    topk_ids: Tensor, num_experts: int, block_m: int
-) -> tuple[Tensor, Tensor, Tensor]:
-    """Sort (token, slot) pairs by expert, padding each expert to a multiple of block_m.
-
-    Stays on the device (no host sync). Returns sorted_ids [max_rows] (padding rows
-    hold the pair count), expert_ids [max_rows / block_m] and num_post_padded [1].
-    """
-    flat = topk_ids.reshape(-1).to(torch.int64)
-    pairs = flat.numel()
-    device = flat.device
-    # At most min(pairs, experts) experts are hit, each padded by < block_m rows.
-    hit_bound = min(pairs, num_experts)
-    max_rows = triton.cdiv(pairs + hit_bound * (block_m - 1), block_m) * block_m
-
-    # bincount syncs with the host on CUDA; scatter_add does not.
-    counts = torch.zeros(num_experts, dtype=torch.int64, device=device)
-    counts.scatter_add_(0, flat, torch.ones_like(flat))
-    padded = (counts + block_m - 1) // block_m * block_m
-    ends = padded.cumsum(0)
-    starts = ends - padded
-    count_starts = counts.cumsum(0) - counts
-    order = torch.argsort(flat, stable=True)
-    sorted_experts = flat[order]
-    rank = torch.arange(pairs, device=device) - count_starts[sorted_experts]
-    sorted_ids = torch.full((max_rows,), pairs, dtype=torch.int32, device=device)
-    sorted_ids[starts[sorted_experts] + rank] = order.to(torch.int32)
-
-    block_starts = torch.arange(0, max_rows, block_m, device=device)
-    expert_ids = torch.searchsorted(ends, block_starts, right=True)
-    expert_ids = expert_ids.clamp_(max=num_experts - 1).to(torch.int32)
-    return sorted_ids, expert_ids, ends[-1:].to(torch.int32)
-
-
-def fused_moe_gemm(
-    a: Tensor,
-    a_scale: Tensor,
-    b: Tensor,
-    b_scale: Tensor,
-    out: Tensor,
-    topk_weights: Tensor | None,
-    sorted_ids: Tensor,
-    expert_ids: Tensor,
-    num_post_padded: Tensor,
-    *,
-    top_k: int,
-    block_m: int,
-    group_m: int,
-) -> None:
-    _, n, k = b.shape
-    grid = (triton.cdiv(sorted_ids.numel(), block_m) * triton.cdiv(n, GEMM_BLOCK_N),)
-    fused_moe_kernel[grid](
-        a,
-        b,
-        out,
-        a_scale,
-        b_scale,
-        # Any valid pointer; the kernel only reads it with MUL_ROUTED_WEIGHT.
-        a_scale if topk_weights is None else topk_weights,
-        sorted_ids,
-        expert_ids,
-        num_post_padded,
-        n,
-        k,
-        a.shape[0] * top_k,
-        a.stride(0),
-        b.stride(0),
-        b.stride(1),
-        b.stride(2),
-        out.stride(0),
-        a_scale.stride(0),
-        b_scale.stride(0),
-        b_scale.stride(1),
-        b_scale.stride(2),
-        GROUP=FP8_BLOCK_SIZE,
-        BLOCK_M=block_m,
-        BLOCK_N=GEMM_BLOCK_N,
-        BLOCK_K=FP8_BLOCK_SIZE,
-        GROUP_M=group_m,
-        TOP_K=top_k,
-        MUL_ROUTED_WEIGHT=topk_weights is not None,
-        num_warps=GEMM_NUM_WARPS,
-        num_stages=GEMM_NUM_STAGES,
-    )
-
-
-def fp8_moe_forward(
-    x: Tensor,
-    topk_ids: Tensor,
-    topk_weights: Tensor,
-    w13_weight: Tensor,
-    w13_scale_inv: Tensor,
-    w2_weight: Tensor,
-    w2_scale_inv: Tensor,
-) -> Tensor:
-    """Route x [N, H] through block-FP8 experts (W8A8) -> weighted expert sum [N, H].
-
-    w13_weight [E, 2I, H] holds gate rows then up rows, w2_weight [E, H, I]; both
-    float8_e4m3fn with fp32 [E, rows / 128, cols / 128] block scales.
-    """
-    tokens, hidden = x.shape
-    if tokens == 0:
-        return torch.zeros_like(x)
-    num_experts, gate_up_rows, _ = w13_weight.shape
-    top_k = topk_ids.shape[-1]
-    pairs = tokens * top_k
-    # Decode routes a handful of rows per expert: smaller M tiles waste less.
-    if pairs <= 64:
-        block_m, group_m = 16, 1
-    elif pairs <= 1024:
-        block_m, group_m = 32, 1
-    else:
-        block_m, group_m = 64, 8
-
-    # Triton launches on the current device; layers may live on another GPU.
-    with torch.cuda.device(x.device):
-        sorted_ids, expert_ids, num_post_padded = moe_align(
-            topk_ids, num_experts, block_m
-        )
-        layout = dict(top_k=top_k, block_m=block_m, group_m=group_m)
-        x_q, x_scale = quantize_groups(quant_kernel, x.contiguous(), hidden)
-        gate_up = torch.empty((pairs, gate_up_rows), dtype=x.dtype, device=x.device)
-        fused_moe_gemm(
-            x_q,
-            x_scale,
-            w13_weight,
-            w13_scale_inv,
-            gate_up,
-            None,
-            sorted_ids,
-            expert_ids,
-            num_post_padded,
-            **layout,
-        )
-        h_q, h_scale = quantize_groups(
-            silu_mul_quant_kernel, gate_up, gate_up_rows // 2
-        )
-        down = torch.empty((pairs, hidden), dtype=torch.float32, device=x.device)
-        fused_moe_gemm(
-            h_q,
-            h_scale,
-            w2_weight,
-            w2_scale_inv,
-            down,
-            topk_weights.reshape(-1).contiguous(),
-            sorted_ids,
-            expert_ids,
-            num_post_padded,
-            top_k=1,
-            block_m=block_m,
-            group_m=group_m,
-        )
-    return down.view(tokens, top_k, hidden).sum(dim=1).to(x.dtype)

fish_speech/models/text2semantic/fish_omni/generate.py

@@ -1,505 +0,0 @@
-import os
-import re
-import time
-from collections.abc import Iterator, Sequence
-from typing import Literal
-
-import torch
-from loguru import logger
-from tokenizers import Tokenizer
-from tqdm import tqdm
-
-from fish_speech.models.text2semantic.errors import InvalidInputError
-from fish_speech.models.text2semantic.fish_omni.config import FishOmniConfig
-from fish_speech.models.text2semantic.fish_omni.model import FishOmniModel
-from fish_speech.models.text2semantic.fish_omni.tower import (
-    RAS_WINDOW_FRAMES,
-    TowerAttention,
-)
-from fish_speech.models.text2semantic.sampling import (
-    DEFAULT_TEMPERATURE,
-    DEFAULT_TOP_K,
-    DEFAULT_TOP_P,
-)
-
-# Special-token strings are assembled with "<" + "|..." so the literal two
-# characters never appear side by side in this source file (tooling mangles them).
-IM_START = "<" + "|im_start|>"
-IM_END = "<" + "|im_end|>"
-VOICE_TAG = "<" + "|voice|>"
-SPEAKER_0 = "<" + "|speaker:0|>"
-
-DRAMA3_TTS_INSTRUCTION = "convert the provided text to speech"
-DRAMA3_REFERENCE_INSTRUCTION = (
-    "convert the provided text to speech reference to the following:"
-)
-DRAMA3_MESSAGE_END = IM_END + "\n"
-DRAMA3_GENERATION_START = IM_START + "assistant\n" + VOICE_TAG
-# Only the prompt builder writes these; other special tokens in caller text are content.
-DRAMA3_STRUCTURAL_TOKENS = (
-    IM_START,
-    IM_END,
-    VOICE_TAG,
-    "<" + "|text|>",
-    "<" + "|interleave|>",
-    "<" + "|semantic:",
-)
-SPEAKER_TAG_PATTERN = re.compile("<[|]" + r"speaker:([0-9]+)" + "[|]>")
-# The tokenizer has one token per speaker, speaker:0 to speaker:15.
-MAX_SPEAKER_ID = 15
-# The Unicode White_Space property; str.strip() without arguments also strips
-# U+001C..U+001F, which would change the empty check and the speaker-tag rule.
-UNICODE_WHITE_SPACE = "".join(
-    map(
-        chr,
-        (*range(0x9, 0xE), 0x20, 0x85, 0xA0, 0x1680, *range(0x2000, 0x200B))
-        + (0x2028, 0x2029, 0x202F, 0x205F, 0x3000),
-    )
-)
-DRAMA3_MAX_PROMPT_ROWS = 8192
-DRAMA3_DEFAULT_MAX_NEW_TOKENS = 2048
-DRAMA3_PREFILL_CHUNK_SIZE = 128
-DRAMA3_MAX_DEFAULT_DEVICES = 4
-# CUDA-graph decode: caches hold the largest default request, with one graph per
-# attended cache length; a step replays the smallest bucket covering its position.
-DRAMA3_GRAPH_CAPACITY = DRAMA3_MAX_PROMPT_ROWS + DRAMA3_DEFAULT_MAX_NEW_TOKENS
-DRAMA3_GRAPH_KV_BUCKETS = (1024, 2048, 4096, DRAMA3_GRAPH_CAPACITY)
-
-FinishReason = Literal["eos", "length"]
-
-
-class Drama3PromptTooLongError(InvalidInputError):
-    """The prompt has more than DRAMA3_MAX_PROMPT_ROWS rows."""
-
-
-def validate_drama3_text(text: str, location: str) -> None:
-    """Non-empty, no structural tokens, and only canonical speaker tags 0..15."""
-    if not text.strip(UNICODE_WHITE_SPACE):
-        raise InvalidInputError(f"{location}: text must not be empty")
-    for token in DRAMA3_STRUCTURAL_TOKENS:
-        if token in text:
-            raise InvalidInputError(
-                f"{location}: {token} is written by the prompt builder"
-            )
-    for marker in re.finditer("<[|]" + "speaker", text):
-        tag = SPEAKER_TAG_PATTERN.match(text, marker.start())
-        digits = tag.group(1) if tag else ""
-        if not digits or str(int(digits)) != digits or int(digits) > MAX_SPEAKER_ID:
-            raise InvalidInputError(
-                f"{location}: speaker tags must be {SPEAKER_0} to "
-                + "<" + f"|speaker:{MAX_SPEAKER_ID}|>; got {text!r}"
-            )
-
-
-def with_speaker_tag(text: str) -> str:
-    """Prefix with speaker tag 0 unless a tag already leads after white space."""
-    if SPEAKER_TAG_PATTERN.match(text.lstrip(UNICODE_WHITE_SPACE)) is None:
-        return SPEAKER_0 + " " + text
-    return text
-
-
-def build_drama3_prompt(
-    text: str,
-    tokenizer: Tokenizer,
-    config: FishOmniConfig,
-    prompt_texts: Sequence[str] = (),
-    prompt_tokens: Sequence[torch.Tensor] = (),
-    history: Sequence[tuple[str, torch.Tensor]] = (),
-) -> torch.Tensor:
-    """CPU int64 rows [rows, 1 + num_codebooks] ending in the open assistant voice turn.
-
-    Text rows are [token, 0, ..., 0]; audio rows are [c0 + semantic_start, c0..cN-1].
-    Every literal is tokenized on its own: tokenizing the joined string changes ids
-    at the seams. history holds earlier (text, codes) chunks, written as completed
-    user and assistant voice turns before the current one.
-    """
-    num_codebooks = config.audio_decoder.num_codebooks
-    codebook_size = config.audio_decoder.vocab_size
-    for token, expected in (
-        ("<" + "|semantic:0|>", config.semantic_start_token_id),
-        ("<" + f"|semantic:{codebook_size - 1}|>", config.semantic_end_token_id),
-        (IM_END, config.eos_token_id),
-    ):
-        if tokenizer.token_to_id(token) != expected:
-            raise ValueError(f"tokenizer/config mismatch for {token}")
-    if len(prompt_texts) != len(prompt_tokens):
-        raise InvalidInputError("every reference needs both a transcript and codes")
-    validate_drama3_text(text, "text")
-
-    def check_codes(codes: torch.Tensor, location: str) -> None:
-        if (
-            codes.dtype.is_floating_point
-            or codes.dtype == torch.bool
-            or codes.ndim != 2
-            or codes.shape[0] != num_codebooks
-            or codes.shape[1] == 0
-        ):
-            raise InvalidInputError(
-                f"{location}: codes must be integer [{num_codebooks}, frames >= 1]; "
-                f"got {codes.dtype} {tuple(codes.shape)}"
-            )
-        if codes.min() < 0 or codes.max() >= codebook_size:
-            raise InvalidInputError(
-                f"{location}: codes must be in [0, {codebook_size})"
-            )
-
-    segments: list[str | torch.Tensor] = [
-        IM_START + "system\n",
-        DRAMA3_REFERENCE_INSTRUCTION if prompt_tokens else DRAMA3_TTS_INSTRUCTION,
-    ]
-
-    for index, (transcript, codes) in enumerate(zip(prompt_texts, prompt_tokens), 1):
-        location = f"reference {index}"
-        validate_drama3_text(transcript, location)
-        check_codes(codes, location)
-        # The tag must follow "Text: " directly; leading white space moves behind it.
-        tagged = with_speaker_tag(transcript)
-        body = tagged.lstrip(UNICODE_WHITE_SPACE)
-        tag = SPEAKER_TAG_PATTERN.match(body)
-        leading = tagged[: len(tagged) - len(body)]
-        transcript = tag.group(0) + leading + body[tag.end() :]
-        segments += [f"\n\n[Reference {index}]\nText: {transcript}\nSpeech: ", codes]
-    segments.append(DRAMA3_MESSAGE_END)
-    for index, (turn_text, codes) in enumerate(history, 1):
-        location = f"history {index}"
-        validate_drama3_text(turn_text, location)
-        check_codes(codes, location)
-        segments += [
-            IM_START + "user\n",
-            with_speaker_tag(turn_text),
-            DRAMA3_MESSAGE_END,
-            DRAMA3_GENERATION_START,
-            codes,
-            DRAMA3_MESSAGE_END,
-        ]
-    segments += [
-        IM_START + "user\n",
-        with_speaker_tag(text),
-        DRAMA3_MESSAGE_END,
-        DRAMA3_GENERATION_START,
-    ]
-
-    rows = []
-    for segment in segments:
-        if isinstance(segment, str):
-            ids = torch.tensor(tokenizer.encode(segment).ids, dtype=torch.long)
-            text_rows = torch.zeros((len(ids), 1 + num_codebooks), dtype=torch.long)
-            text_rows[:, 0] = ids
-            rows.append(text_rows)
-        else:
-            codes = segment.to(device="cpu", dtype=torch.long)
-            main = codes[:1] + config.semantic_start_token_id
-            rows.append(torch.cat((main, codes)).T)
-    rows = torch.cat(rows)
-    if len(rows) > DRAMA3_MAX_PROMPT_ROWS:
-        raise Drama3PromptTooLongError(
-            f"prompt has {len(rows)} rows; at most {DRAMA3_MAX_PROMPT_ROWS}"
-        )
-    return rows
-
-
-@torch.inference_mode()
-def generate_drama3(
-    model: FishOmniModel,
-    prompt: torch.Tensor,
-    *,
-    max_new_tokens: int,
-    temperature: float = DEFAULT_TEMPERATURE,
-    top_k: int = DEFAULT_TOP_K,
-    top_p: float = DEFAULT_TOP_P,
-    prefill_chunk_size: int = DRAMA3_PREFILL_CHUNK_SIZE,
-) -> tuple[torch.Tensor, FinishReason]:
-    """Sample frames after the prompt rows; returns CPU [frames, 1 + N] without EOS.
-
-    The tower cache of prompt row p is filled with frame p + 1 as its label; the
-    last prompt row has none yet, so the first decode step writes that slot.
-    """
-    # The sampler keeps the top-ranked token whatever top_k is, so top_k < 1 would
-    # silently turn into greedy decoding.
-    if top_k < 1:
-        raise InvalidInputError(f"top_k must be at least 1; got {top_k}")
-    backbone, tower = model.backbone, model.audio_decoder
-    prompt_rows = len(prompt)
-    capacity = prompt_rows + max_new_tokens
-    if max_new_tokens < 1 or capacity > model.config.text.max_position_embeddings:
-        raise ValueError(
-            f"max_new_tokens={max_new_tokens} does not fit after {prompt_rows} "
-            "prompt rows"
-        )
-    graphs = Drama3DecodeGraphs.for_model(model, capacity)
-    if graphs is None:
-        ensure_drama3_caches(model, capacity)
-    frames = torch.zeros(
-        (capacity, prompt.shape[1]), dtype=torch.long, device=tower.device
-    )
-    frames[:prompt_rows] = prompt.to(tower.device)
-    positions = torch.arange(capacity, device=tower.device)
-
-    started = time.perf_counter()
-    for start in range(0, prompt_rows, prefill_chunk_size):
-        end = min(start + prefill_chunk_size, prompt_rows)
-        chunk = frames[start:end]
-        embeddings = tower.embed_backbone(chunk, backbone.embed(chunk[:, 0]))
-        hidden = backbone(
-            inputs_embeds=embeddings[None], input_pos=positions[start:end]
-        )
-        hidden = hidden[0].to(tower.device)
-        paired_end = min(end, prompt_rows - 1)
-        if paired_end > start:
-            tower.prefill(
-                hidden[: paired_end - start],
-                positions[start:paired_end],
-                frames[start:paired_end],
-                frames[start + 1 : paired_end + 1],
-            )
-    logger.info(
-        f"Drama 3 prefill: {prompt_rows} rows in {time.perf_counter() - started:.2f}s"
-    )
-
-    hidden = hidden[-1:]
-    produced, finish_reason = 0, "length"
-    if graphs is not None:
-        graphs.start(
-            hidden, prompt_rows - 1, frames[:prompt_rows], temperature, top_k, top_p
-        )
-    for step in tqdm(range(max_new_tokens)):
-        position = prompt_rows - 1 + step
-        if graphs is not None:
-            sampled, main = graphs.step(position)
-        else:
-            logits = backbone.compute_logits(hidden).to(tower.device)
-            sampled = tower.decode_step(
-                logits,
-                hidden,
-                positions[position : position + 1],
-                frames[position : position + 1],
-                frames[: position + 1, 0],
-                temperature=temperature,
-                top_k=top_k,
-                top_p=top_p,
-            )
-            main = sampled[0, 0].item()
-        if main == model.config.eos_token_id:
-            finish_reason = "eos"
-            break
-        frames[position + 1] = sampled[0]
-        produced += 1
-        if graphs is None and step + 1 < max_new_tokens:
-            embeddings = tower.embed_backbone(sampled, backbone.embed(sampled[:, 0]))
-            hidden = backbone(
-                inputs_embeds=embeddings[None],
-                input_pos=positions[position + 1 : position + 2],
-            )[0].to(tower.device)
-    return frames[prompt_rows : prompt_rows + produced].cpu(), finish_reason
-
-
-def ensure_drama3_caches(model: FishOmniModel, capacity: int) -> None:
-    """Grow the backbone and tower caches to capacity rows, reused across requests.
-
-    Stale rows past a request's positions are masked out.
-    """
-    if model.cache_capacity >= capacity:
-        return
-    if model.decode_graphs is not None:
-        # The graphs captured the old cache buffers.
-        model.decode_graphs = None
-        torch.cuda.empty_cache()
-    model.backbone.setup_caches(max_batch_size=1, max_seq_len=capacity)
-    model.audio_decoder.setup_caches(capacity)
-    model.cache_capacity = capacity
-
-
-class Drama3DecodeGraphs:
-    """CUDA graphs of one decode step: lm_head -> tower (sample a frame) -> backbone.
-
-    Eager decode is bound by kernel launches (thousands of small kernels per step).
-    Captured once for a model on a single CUDA device; FISH_DRAMA3_CUDA_GRAPH=0 disables.
-    """
-
-    def __init__(self, model: FishOmniModel) -> None:
-        backbone, tower = model.backbone, model.audio_decoder
-        self.model = model
-        device = tower.device
-        width = 1 + model.config.audio_decoder.num_codebooks
-        hidden_size = model.config.text.hidden_size
-        self.hidden = torch.zeros((1, hidden_size), dtype=backbone.dtype, device=device)
-        self.next_hidden = torch.zeros_like(self.hidden)
-        self.position = torch.zeros(1, dtype=torch.long, device=device)
-        self.frame = torch.zeros((1, width), dtype=torch.long, device=device)
-        self.sampled = torch.zeros_like(self.frame)
-        self.history = torch.full(
-            (RAS_WINDOW_FRAMES,), -1, dtype=torch.long, device=device
-        )
-        self.temperature = torch.zeros((), dtype=torch.float32, device=device)
-        self.top_k = torch.zeros((), dtype=torch.long, device=device)
-        self.top_p = torch.zeros((), dtype=torch.float32, device=device)
-        self.finite = torch.ones((), dtype=torch.bool, device=device)
-        self.graphs: dict[int, torch.cuda.CUDAGraph] = {}
-        self.pool = torch.cuda.graph_pool_handle()
-        capacity = self.graph_capacity(model)
-        self.buckets = tuple(
-            sorted({min(b, capacity) for b in DRAMA3_GRAPH_KV_BUCKETS})
-        )
-
-    @staticmethod
-    def graph_capacity(model: FishOmniModel) -> int:
-        return min(DRAMA3_GRAPH_CAPACITY, model.config.text.max_position_embeddings)
-
-    @classmethod
-    def for_model(
-        cls, model: FishOmniModel, capacity: int
-    ) -> "Drama3DecodeGraphs | None":
-        """The model's graphs (captured on first use), or None to decode eagerly."""
-        if os.environ.get("FISH_DRAMA3_CUDA_GRAPH", "1") == "0":
-            return None
-        if capacity > cls.graph_capacity(model):
-            return None
-        devices = {p.device for p in model.backbone.parameters()}
-        devices.add(model.audio_decoder.device)
-        if len(devices) != 1 or next(iter(devices)).type != "cuda":
-            return None
-        if model.decode_graphs is None:
-            # Captured before the first request's prefill, which resets the caches.
-            ensure_drama3_caches(model, cls.graph_capacity(model))
-            graphs = cls(model)
-            started = time.perf_counter()
-            with torch.cuda.device(model.audio_decoder.device):
-                for kv_len in graphs.buckets:
-                    graphs.capture(kv_len)
-            logger.info(
-                f"Drama 3 decode CUDA graphs: kv buckets {graphs.buckets} captured in "
-                f"{time.perf_counter() - started:.1f}s"
-            )
-            model.decode_graphs = graphs
-        return model.decode_graphs
-
-    def tower_attentions(self) -> Iterator[TowerAttention]:
-        tower = self.model.audio_decoder
-        for layer in (*tower.shared_layers, *tower.residual_layers):
-            yield layer.transformer_layer.attention
-
-    def run(self, kv_len: int) -> None:
-        backbone, tower = self.model.backbone, self.model.audio_decoder
-        logits = backbone.compute_logits(self.hidden)
-        self.finite.copy_(torch.isfinite(logits).all())
-        sampled = tower.decode_step(
-            logits,
-            self.hidden,
-            self.position,
-            self.frame,
-            self.history,
-            temperature=self.temperature,
-            top_k=self.top_k,
-            top_p=self.top_p,
-            check_finite=False,
-        )
-        self.sampled.copy_(sampled)
-        embeddings = tower.embed_backbone(sampled, backbone.embed(sampled[:, 0]))
-        hidden = backbone(
-            inputs_embeds=embeddings[None],
-            input_pos=self.position + 1,
-            static_kv_len=kv_len,
-        )
-        self.next_hidden.copy_(hidden[0])
-
-    def capture(self, kv_len: int) -> None:
-        for attention in self.tower_attentions():
-            attention.static_kv_len = kv_len
-        try:
-            stream = torch.cuda.Stream()
-            stream.wait_stream(torch.cuda.current_stream())
-            with torch.cuda.stream(stream):
-                for _ in range(2):
-                    self.run(kv_len)
-            torch.cuda.current_stream().wait_stream(stream)
-            graph = torch.cuda.CUDAGraph()
-            with torch.cuda.graph(graph, pool=self.pool):
-                self.run(kv_len)
-            self.graphs[kv_len] = graph
-        finally:
-            for attention in self.tower_attentions():
-                attention.static_kv_len = None
-
-    def start(
-        self,
-        hidden: torch.Tensor,
-        position: int,
-        prompt_frames: torch.Tensor,
-        temperature: float,
-        top_k: int,
-        top_p: float,
-    ) -> None:
-        """Load the state after prefill: last hidden row, its position and frame."""
-        self.hidden.copy_(hidden)
-        self.position.fill_(position)
-        self.frame.copy_(prompt_frames[-1:])
-        window = prompt_frames[-RAS_WINDOW_FRAMES:, 0]
-        self.history.fill_(-1)
-        self.history[RAS_WINDOW_FRAMES - len(window) :].copy_(window)
-        self.temperature.fill_(temperature)
-        self.top_k.fill_(top_k)
-        self.top_p.fill_(top_p)
-
-    def step(self, position: int) -> tuple[torch.Tensor, int]:
-        """Sample the frame after position and run the backbone on it.
-
-        Returns the frame [1, 1 + N] on the device and its main token.
-        """
-        kv_len = next(b for b in self.buckets if b >= position + 2)
-        self.graphs[kv_len].replay()
-        sampled = self.sampled.clone()
-        status = torch.stack((sampled[0, 0], self.finite.long()))
-        # Inputs of the next step, queued before the host waits on status.
-        self.hidden.copy_(self.next_hidden)
-        self.frame.copy_(sampled)
-        self.history.copy_(torch.cat((self.history[1:], sampled[0, :1])))
-        self.position.add_(1)
-        main, finite = status.tolist()
-        if not finite:
-            raise FloatingPointError(
-                "non-finite logits; refusing to sample invalid audio"
-            )
-        return sampled, main
-
-
-def resolve_drama3_devices(device: str) -> tuple[list[str], str]:
-    """Split "cuda:0,cuda:1,..." into backbone devices and the tower (+ codec) device.
-
-    The last device holds the tower; a single device holds everything. Plain "cuda"
-    takes the first min(4, visible) GPUs.
-    """
-    specs = [spec.strip() for spec in device.split(",") if spec.strip()]
-    if not specs:
-        raise ValueError("no device given")
-    if len(specs) == 1 and specs[0] == "cuda":
-        count = min(torch.cuda.device_count(), DRAMA3_MAX_DEFAULT_DEVICES)
-        if count == 0:
-            raise RuntimeError("Drama 3 inference needs CUDA, but no GPU is visible")
-        specs = [f"cuda:{index}" for index in range(count)]
-    for spec in specs:
-        parsed = torch.device(spec)
-        if (
-            parsed.type != "cuda"
-            or parsed.index is None
-            or parsed.index >= torch.cuda.device_count()
-        ):
-            raise ValueError(f"Drama 3 needs available cuda:N devices; got {spec!r}")
-    if len(set(specs)) != len(specs):
-        raise ValueError(f"duplicate devices in {device!r}")
-    if len(specs) == 1:
-        return specs, specs[0]
-    return specs[:-1], specs[-1]
-
-
-def resolve_drama3_dtype(precision: torch.dtype, devices: list[str]) -> torch.dtype:
-    """bfloat16 needs SM80+; on older GPUs (V100) Drama 3 runs in float16 instead."""
-    if precision not in (torch.float16, torch.bfloat16):
-        raise ValueError(f"Drama 3 runs in float16 or bfloat16; got {precision}")
-    if precision == torch.bfloat16 and any(
-        torch.cuda.get_device_capability(device) < (8, 0) for device in devices
-    ):
-        logger.warning(
-            "bfloat16 needs SM80+; running Drama 3 in float16 with an fp32 residual stream"
-        )
-        return torch.float16
-    return precision

fish_speech/models/text2semantic/fish_omni/layers.py

@@ -1,410 +0,0 @@
-from dataclasses import dataclass
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch import Tensor
-from torch.types import Device
-
-from fish_speech.models.text2semantic.fish_omni.config import FishOmniTextConfig
-
-GDN_CHUNK_SIZE = 64
-
-
-class ZeroCenteredRMSNorm(nn.Module):
-    """RMSNorm scaling by 1 + weight, computed in fp32."""
-
-    def __init__(
-        self,
-        dim: int,
-        eps: float,
-        *,
-        dtype: torch.dtype | None = None,
-        device: Device = None,
-    ) -> None:
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.zeros(dim, dtype=dtype, device=device))
-
-    def forward(self, x: Tensor) -> Tensor:
-        out = x.float()
-        out = out * torch.rsqrt(out.pow(2).mean(-1, keepdim=True) + self.eps)
-        return (out * (1.0 + self.weight.float())).type_as(x)
-
-
-class GatedRMSNorm(nn.Module):
-    """Gated DeltaNet output norm: weight * rmsnorm(x) * silu(gate)."""
-
-    def __init__(
-        self,
-        dim: int,
-        eps: float,
-        *,
-        dtype: torch.dtype | None = None,
-        device: Device = None,
-    ) -> None:
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim, dtype=dtype, device=device))
-
-    def forward(self, x: Tensor, gate: Tensor) -> Tensor:
-        dtype = x.dtype
-        out = x.float()
-        out = out * torch.rsqrt(out.pow(2).mean(-1, keepdim=True) + self.eps)
-        out = self.weight * out.to(dtype)
-        return (out * F.silu(gate.float())).to(dtype)
-
-
-def apply_neox_rope(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
-    """Rotate the first rotary_dim channels of x [B, T, heads, head_dim] (half-split)."""
-    rotary_dim = cos.shape[-1]
-    rotary, passthrough = x[..., :rotary_dim].float(), x[..., rotary_dim:]
-    first, second = rotary.chunk(2, dim=-1)
-    rotated = torch.cat((-second, first), dim=-1)
-    cos, sin = cos[None, :, None, :], sin[None, :, None, :]
-    out = rotary * cos + rotated * sin
-    return torch.cat((out.to(x.dtype), passthrough), dim=-1)
-
-
-def grouped_attention(
-    query: Tensor,
-    key: Tensor,
-    value: Tensor,
-    mask: Tensor | None,
-    *,
-    is_causal: bool = False,
-    scale: float | None = None,
-) -> Tensor:
-    """GQA SDPA: query [B, H, T, D], key / value [B, KVH, S, D], mask [T, S]."""
-    batch, heads, length, head_dim = query.shape
-    kv_heads = key.shape[1]
-    groups = heads // kv_heads
-    if groups > 1 and length == 1:
-        # Decode: fold each KV head's query heads into rows, so the cache is read
-        # once; the [1, S] mask broadcasts over them.
-        folded = query.reshape(batch, kv_heads, groups, head_dim)
-        out = F.scaled_dot_product_attention(
-            folded, key, value, attn_mask=mask, scale=scale
-        )
-        return out.reshape(batch, heads, 1, head_dim)
-    if groups > 1:
-        key = key.repeat_interleave(groups, dim=1)
-        value = value.repeat_interleave(groups, dim=1)
-    return F.scaled_dot_product_attention(
-        query, key, value, attn_mask=mask, is_causal=is_causal, scale=scale
-    )
-
-
-@dataclass(kw_only=True)
-class LayerContext:
-    """Per-forward inputs shared by the layers of one device."""
-
-    input_pos: Tensor  # [T] contiguous positions of the rows
-    cos: Tensor
-    sin: Tensor
-    kv_len: int  # attended cache prefix (last position + 1)
-    continue_state: bool  # the rows continue a cached sequence
-    use_cache: bool
-
-
-class FullAttention(nn.Module):
-    """Gated softmax attention: GQA, q/k norm, partial RoPE, sigmoid output gate."""
-
-    def __init__(
-        self,
-        config: FishOmniTextConfig,
-        *,
-        dtype: torch.dtype | None = None,
-        device: Device = None,
-    ) -> None:
-        super().__init__()
-        self.num_heads = config.num_attention_heads
-        self.num_kv_heads = config.num_key_value_heads
-        self.head_dim = config.head_dim
-        hidden = config.hidden_size
-        linear = dict(bias=False, dtype=dtype, device=device)
-        # Output rows: each head's [query, gate] pair.
-        self.q_proj = nn.Linear(hidden, self.num_heads * self.head_dim * 2, **linear)
-        self.k_proj = nn.Linear(hidden, self.num_kv_heads * self.head_dim, **linear)
-        self.v_proj = nn.Linear(hidden, self.num_kv_heads * self.head_dim, **linear)
-        self.o_proj = nn.Linear(self.num_heads * self.head_dim, hidden, **linear)
-        norm = dict(dtype=dtype, device=device)
-        self.q_norm = ZeroCenteredRMSNorm(self.head_dim, config.rms_norm_eps, **norm)
-        self.k_norm = ZeroCenteredRMSNorm(self.head_dim, config.rms_norm_eps, **norm)
-        self.k_cache: Tensor | None = None
-        self.v_cache: Tensor | None = None
-
-    def allocate_cache(
-        self, batch_size: int, max_seq_len: int, dtype: torch.dtype, device: Device
-    ) -> None:
-        shape = (batch_size, self.num_kv_heads, max_seq_len, self.head_dim)
-        self.k_cache = torch.zeros(shape, dtype=dtype, device=device)
-        self.v_cache = torch.zeros(shape, dtype=dtype, device=device)
-
-    def forward(self, x: Tensor, context: LayerContext) -> Tensor:
-        batch, length, _ = x.shape
-        query, gate = (
-            self.q_proj(x)
-            .view(batch, length, self.num_heads, 2 * self.head_dim)
-            .chunk(2, dim=-1)
-        )
-        gate = gate.reshape(batch, length, self.num_heads * self.head_dim)
-        query = self.q_norm(query)
-        key = self.k_proj(x).view(batch, length, self.num_kv_heads, self.head_dim)
-        key = self.k_norm(key)
-        value = self.v_proj(x).view(batch, length, self.num_kv_heads, self.head_dim)
-        query = apply_neox_rope(query, context.cos, context.sin).transpose(1, 2)
-        key = apply_neox_rope(key, context.cos, context.sin).transpose(1, 2)
-        value = value.transpose(1, 2)
-
-        if context.use_cache:
-            assert self.k_cache is not None and self.v_cache is not None
-            self.k_cache[:batch].index_copy_(2, context.input_pos, key)
-            self.v_cache[:batch].index_copy_(2, context.input_pos, value)
-            key = self.k_cache[:batch, :, : context.kv_len]
-            value = self.v_cache[:batch, :, : context.kv_len]
-            key_positions = torch.arange(context.kv_len, device=x.device)
-            mask = key_positions[None, :] <= context.input_pos[:, None]
-            is_causal = False
-        else:
-            mask = None
-            is_causal = length > 1
-
-        attended = grouped_attention(
-            query, key, value, mask, is_causal=is_causal, scale=self.head_dim**-0.5
-        )
-        attended = attended.transpose(1, 2).reshape(
-            batch, length, self.num_heads * self.head_dim
-        )
-        return self.o_proj(attended * torch.sigmoid(gate))
-
-
-def l2norm(x: Tensor, eps: float = 1e-6) -> Tensor:
-    return x * torch.rsqrt((x * x).sum(dim=-1, keepdim=True) + eps)
-
-
-def chunk_gated_delta_rule(
-    query: Tensor,
-    key: Tensor,
-    value: Tensor,
-    decay: Tensor,
-    beta: Tensor,
-    initial_state: Tensor | None = None,
-    chunk_size: int = GDN_CHUNK_SIZE,
-) -> tuple[Tensor, Tensor]:
-    """Gated delta rule over whole chunks, in fp32.
-
-    query / key [B, T, H, dk] (not yet l2-normalized), value [B, T, H, dv], decay
-    (log space, <= 0) and beta [B, T, H], initial_state [B, H, dk, dv]. Returns the
-    output [B, T, H, dv] in the query dtype and the final fp32 state.
-    """
-    output_dtype = query.dtype
-    batch, length, heads, key_dim = key.shape
-    value_dim = value.shape[-1]
-    query, key, value, beta, decay = (
-        t.transpose(1, 2).to(torch.float32).contiguous()
-        for t in (query, key, value, beta, decay)
-    )
-    query = l2norm(query) * key_dim**-0.5
-    key = l2norm(key)
-
-    pad = (chunk_size - length % chunk_size) % chunk_size
-    query, key, value = (F.pad(t, (0, 0, 0, pad)) for t in (query, key, value))
-    beta, decay = (F.pad(t, (0, pad)) for t in (beta, decay))
-    chunk_count = (length + pad) // chunk_size
-
-    value_beta = value * beta.unsqueeze(-1)
-    key_beta = key * beta.unsqueeze(-1)
-    query, key, key_beta, value_beta = (
-        t.reshape(batch, heads, chunk_count, chunk_size, t.shape[-1])
-        for t in (query, key, key_beta, value_beta)
-    )
-    decay = decay.reshape(batch, heads, chunk_count, chunk_size)
-    upper = torch.ones(
-        chunk_size, chunk_size, dtype=torch.bool, device=query.device
-    ).triu(1)
-    cumulative = decay.cumsum(dim=-1)
-    pairwise = (
-        (cumulative.unsqueeze(-1) - cumulative.unsqueeze(-2))
-        .masked_fill(upper, float("-inf"))
-        .exp()
-    )
-
-    system = (key_beta @ key.transpose(-1, -2)) * pairwise
-    intra = (query @ key.transpose(-1, -2)) * pairwise
-    decayed_key_beta = key_beta * cumulative.exp().unsqueeze(-1)
-    new_values = torch.linalg.solve_triangular(
-        system, value_beta, upper=False, unitriangular=True
-    )
-    key_cumulative_decay = torch.linalg.solve_triangular(
-        system, decayed_key_beta, upper=False, unitriangular=True
-    )
-
-    if initial_state is None:
-        state = query.new_zeros(batch, heads, key_dim, value_dim)
-    else:
-        state = initial_state.to(torch.float32)
-    output = torch.empty_like(new_values)
-    query = query * cumulative.exp().unsqueeze(-1)
-    key = key * (cumulative[..., -1:] - cumulative).exp().unsqueeze(-1)
-    chunk_decay = cumulative[..., -1].exp()[..., None, None]
-    for index in range(chunk_count):
-        corrected = new_values[:, :, index] - key_cumulative_decay[:, :, index] @ state
-        output[:, :, index] = (
-            query[:, :, index] @ state + intra[:, :, index] @ corrected
-        )
-        state = (
-            state * chunk_decay[:, :, index]
-            + key[:, :, index].transpose(-1, -2) @ corrected
-        )
-
-    output = output.reshape(batch, heads, -1, value_dim)[:, :, :length]
-    return output.transpose(1, 2).contiguous().to(output_dtype), state
-
-
-def recurrent_gated_delta_rule(
-    query: Tensor,
-    key: Tensor,
-    value: Tensor,
-    decay: Tensor,
-    beta: Tensor,
-    initial_state: Tensor | None = None,
-) -> tuple[Tensor, Tensor]:
-    """Token-by-token gated delta rule; same contract as chunk_gated_delta_rule."""
-    output_dtype = query.dtype
-    batch, length, heads, key_dim = key.shape
-    value_dim = value.shape[-1]
-    query, key, value, beta, decay = (
-        t.transpose(1, 2).to(torch.float32).contiguous()
-        for t in (query, key, value, beta, decay)
-    )
-    query = l2norm(query) * key_dim**-0.5
-    key = l2norm(key)
-    if initial_state is None:
-        state = value.new_zeros(batch, heads, key_dim, value_dim)
-    else:
-        state = initial_state.to(torch.float32)
-    output = torch.empty_like(value)
-    for index in range(length):
-        state = state * decay[:, :, index].exp()[..., None, None]
-        k_t = key[:, :, index]
-        remembered = (state * k_t.unsqueeze(-1)).sum(dim=-2)
-        delta = (value[:, :, index] - remembered) * beta[:, :, index].unsqueeze(-1)
-        state = state + k_t.unsqueeze(-1) * delta.unsqueeze(-2)
-        output[:, :, index] = (state * query[:, :, index].unsqueeze(-1)).sum(dim=-2)
-    return output.transpose(1, 2).contiguous().to(output_dtype), state
-
-
-class GatedDeltaNet(nn.Module):
-    def __init__(
-        self,
-        config: FishOmniTextConfig,
-        *,
-        dtype: torch.dtype | None = None,
-        device: Device = None,
-    ) -> None:
-        super().__init__()
-        hidden = config.hidden_size
-        self.num_k_heads = config.linear_num_key_heads
-        self.num_v_heads = config.linear_num_value_heads
-        self.head_k_dim = config.linear_key_head_dim
-        self.head_v_dim = config.linear_value_head_dim
-        self.key_dim = self.num_k_heads * self.head_k_dim
-        self.value_dim = self.num_v_heads * self.head_v_dim
-        self.conv_dim = 2 * self.key_dim + self.value_dim
-        self.conv_kernel_size = config.linear_conv_kernel_dim
-        linear = dict(bias=False, dtype=dtype, device=device)
-        # Output rows: query, key, value.
-        self.in_proj_qkv = nn.Linear(hidden, self.conv_dim, **linear)
-        self.in_proj_z = nn.Linear(hidden, self.value_dim, **linear)
-        self.in_proj_b = nn.Linear(hidden, self.num_v_heads, **linear)
-        self.in_proj_a = nn.Linear(hidden, self.num_v_heads, **linear)
-        self.conv1d = nn.Conv1d(
-            self.conv_dim,
-            self.conv_dim,
-            kernel_size=self.conv_kernel_size,
-            groups=self.conv_dim,
-            bias=False,
-            dtype=dtype,
-            device=device,
-        )
-        self.dt_bias = nn.Parameter(
-            torch.ones(self.num_v_heads, dtype=dtype, device=device)
-        )
-        self.A_log = nn.Parameter(
-            torch.zeros(self.num_v_heads, dtype=dtype, device=device)
-        )
-        self.norm = GatedRMSNorm(
-            self.head_v_dim, config.rms_norm_eps, dtype=dtype, device=device
-        )
-        self.out_proj = nn.Linear(self.value_dim, hidden, **linear)
-        self.conv_state: Tensor | None = None
-        self.recurrent_state: Tensor | None = None
-
-    def allocate_cache(
-        self, batch_size: int, dtype: torch.dtype, device: Device
-    ) -> None:
-        self.conv_state = torch.zeros(
-            batch_size,
-            self.conv_dim,
-            self.conv_kernel_size - 1,
-            dtype=dtype,
-            device=device,
-        )
-        self.recurrent_state = torch.zeros(
-            batch_size,
-            self.num_v_heads,
-            self.head_k_dim,
-            self.head_v_dim,
-            dtype=torch.float32,
-            device=device,
-        )
-
-    def forward(self, x: Tensor, context: LayerContext) -> Tensor:
-        batch, length, _ = x.shape
-        mixed = self.in_proj_qkv(x).transpose(1, 2)  # [B, C, T]
-        history = self.conv_kernel_size - 1
-        if context.continue_state:
-            assert self.conv_state is not None
-            previous = self.conv_state[:batch].to(mixed.dtype)
-        else:
-            previous = mixed.new_zeros(batch, self.conv_dim, history)
-        padded = torch.cat((previous, mixed), dim=-1)
-        if context.use_cache:
-            assert self.conv_state is not None
-            self.conv_state[:batch].copy_(padded[..., -history:])
-        mixed = F.silu(F.conv1d(padded, self.conv1d.weight, groups=self.conv_dim))
-        query, key, value = mixed.transpose(1, 2).split(
-            [self.key_dim, self.key_dim, self.value_dim], dim=-1
-        )
-        query = query.reshape(batch, length, self.num_k_heads, self.head_k_dim)
-        key = key.reshape(batch, length, self.num_k_heads, self.head_k_dim)
-        value = value.reshape(batch, length, self.num_v_heads, self.head_v_dim)
-        z = self.in_proj_z(x).reshape(batch, length, self.num_v_heads, self.head_v_dim)
-        beta = torch.sigmoid(self.in_proj_b(x).float())
-        decay = -self.A_log.float().exp() * F.softplus(
-            self.in_proj_a(x).float() + self.dt_bias.float()
-        )
-        groups = self.num_v_heads // self.num_k_heads
-        if groups > 1:
-            query = query.repeat_interleave(groups, dim=2)
-            key = key.repeat_interleave(groups, dim=2)
-
-        initial_state = None
-        if context.continue_state:
-            assert self.recurrent_state is not None
-            initial_state = self.recurrent_state[:batch]
-        rule = (
-            recurrent_gated_delta_rule
-            if context.continue_state and length == 1
-            else chunk_gated_delta_rule
-        )
-        core, final_state = rule(query, key, value, decay, beta, initial_state)
-        if context.use_cache:
-            assert self.recurrent_state is not None
-            self.recurrent_state[:batch].copy_(final_state)
-
-        core = self.norm(
-            core.reshape(-1, self.head_v_dim), z.reshape(-1, self.head_v_dim)
-        ).reshape(batch, length, self.value_dim)
-        return self.out_proj(core)

fish_speech/models/text2semantic/fish_omni/model.py

@@ -1,147 +0,0 @@
-import json
-import os
-from collections.abc import Iterator, Sequence
-from pathlib import Path
-from typing import TYPE_CHECKING
-
-import torch
-import torch.nn as nn
-from loguru import logger
-from safetensors import safe_open
-from tokenizers import Tokenizer
-from torch import Tensor
-from torch.types import Device
-
-from fish_speech.models.text2semantic.fish_omni.backbone import (
-    CHECKPOINT_BACKBONE_PREFIX,
-    CHECKPOINT_LM_HEAD_KEY,
-    FishOmniBackbone,
-    load_backbone_weights,
-)
-from fish_speech.models.text2semantic.fish_omni.config import (
-    FISH_OMNI_MODEL_TYPE,
-    FP8_QUANT_METHOD,
-    FishOmniConfig,
-)
-from fish_speech.models.text2semantic.fish_omni.fp8_moe import fp8_gemm_supported
-from fish_speech.models.text2semantic.fish_omni.tower import (
-    FishOmniAudioDecoder,
-    load_audio_decoder_weights,
-)
-
-if TYPE_CHECKING:
-    from fish_speech.models.text2semantic.fish_omni.generate import Drama3DecodeGraphs
-
-AUDIO_DECODER_PREFIXES = (
-    "audio_decoder.",
-    "audio_decoder_mtp_shared.",
-    "audio_decoder_mtp.",
-)
-
-
-class FishOmniModel(nn.Module):
-    """Drama 3 TTS model: backbone (over one or more devices) + codebook tower + tokenizer."""
-
-    def __init__(
-        self,
-        config: FishOmniConfig,
-        backbone: FishOmniBackbone,
-        audio_decoder: FishOmniAudioDecoder,
-        tokenizer: Tokenizer,
-    ) -> None:
-        super().__init__()
-        self.config = config
-        self.backbone = backbone
-        self.audio_decoder = audio_decoder
-        self.tokenizer = tokenizer
-        # Decode state reused across requests; managed by fish_omni.generate.
-        self.cache_capacity = 0
-        self.decode_graphs: "Drama3DecodeGraphs | None" = None
-
-    @property
-    def dtype(self) -> torch.dtype:
-        return self.backbone.dtype
-
-    @property
-    def device(self) -> torch.device:
-        """Where frames are sampled and codes come out (the tower's device)."""
-        return self.audio_decoder.device
-
-    @classmethod
-    def from_pretrained(
-        cls,
-        checkpoint_dir: str | Path,
-        *,
-        devices: Sequence[Device],
-        tower_device: Device,
-        dtype: torch.dtype,
-    ) -> "FishOmniModel":
-        """Load backbone (spread over devices), tower and tokenizer.
-
-        FP8 checkpoint experts stay FP8 when every backbone device is SM89+ and are
-        dequantized to dtype otherwise; env FISH_DRAMA3_FP8_EXPERTS=0/1 forces either way.
-        """
-        checkpoint_dir = Path(checkpoint_dir)
-        with open(checkpoint_dir / "config.json", encoding="utf-8") as f:
-            config = FishOmniConfig.from_dict(json.load(f))
-        fp8_available = config.fp8_experts and fp8_gemm_supported(devices)
-        forced = os.environ.get("FISH_DRAMA3_FP8_EXPERTS", "").strip()
-        fp8_experts = forced == "1" if forced in ("0", "1") else fp8_available
-        if fp8_experts and not fp8_available:
-            raise ValueError(
-                f"FP8 experts need a {FP8_QUANT_METHOD} checkpoint and SM89+ CUDA "
-                f"devices with Triton; got {list(devices)}"
-            )
-        logger.info(
-            "FishOmni experts: "
-            + ("block-FP8 W8A8" if fp8_experts else f"dequantized to {dtype}")
-        )
-
-        with torch.device("meta"):
-            backbone = FishOmniBackbone(
-                config.text, fp8_experts=fp8_experts, dtype=dtype
-            )
-            audio_decoder = FishOmniAudioDecoder(config, dtype=dtype)
-        backbone.materialize(devices)
-        load_backbone_weights(
-            backbone,
-            iter_checkpoint_tensors(
-                checkpoint_dir, skip_prefixes=AUDIO_DECODER_PREFIXES
-            ),
-        )
-        audio_decoder.to_empty(device=tower_device)
-        load_audio_decoder_weights(
-            audio_decoder,
-            iter_checkpoint_tensors(
-                checkpoint_dir,
-                skip_prefixes=(CHECKPOINT_BACKBONE_PREFIX, CHECKPOINT_LM_HEAD_KEY),
-            ),
-        )
-        tokenizer = Tokenizer.from_file(str(checkpoint_dir / "tokenizer.json"))
-        return cls(config, backbone, audio_decoder, tokenizer).eval()
-
-
-def iter_checkpoint_tensors(
-    checkpoint_dir: Path, *, skip_prefixes: tuple[str, ...]
-) -> Iterator[tuple[str, Tensor]]:
-    """(name, CPU tensor) of every safetensors tensor not under skip_prefixes."""
-    index_path = checkpoint_dir / "model.safetensors.index.json"
-    if index_path.exists():
-        with open(index_path, encoding="utf-8") as f:
-            weight_map = json.load(f)["weight_map"]
-        shards = sorted({checkpoint_dir / name for name in weight_map.values()})
-    else:
-        shards = [checkpoint_dir / "model.safetensors"]
-    for shard in shards:
-        with safe_open(str(shard), framework="pt", device="cpu") as f:
-            for key in f.keys():
-                if not key.startswith(skip_prefixes):
-                    yield key, f.get_tensor(key)
-
-
-def is_fish_omni_checkpoint(path: str | Path) -> bool:
-    config_path = Path(path) / "config.json"
-    if not config_path.is_file():
-        return False
-    with open(config_path, encoding="utf-8") as f:
-        return json.load(f).get("model_type") == FISH_OMNI_MODEL_TYPE

fish_speech/models/text2semantic/fish_omni/moe.py

@@ -1,223 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from loguru import logger
-from torch import Tensor
-from torch.types import Device
-
-from fish_speech.models.text2semantic.fish_omni.config import FishOmniTextConfig
-from fish_speech.models.text2semantic.fish_omni.fp8_moe import (
-    FP8_BLOCK_SIZE,
-    FP8_DTYPE,
-    fp8_moe_forward,
-)
-
-# Up to this many (token, expert) pairs run every expert in one batched matmul
-# (decode); larger batches loop over the experts they hit (prefill).
-GATHERED_EXPERT_MAX_ROWS = 64
-FLOAT16_MAX = torch.finfo(torch.float16).max
-
-
-def dequantize_fp8_blocks(
-    weight: Tensor, scale_inv: Tensor, dtype: torch.dtype
-) -> Tensor:
-    """Block-FP8 [..., out, in] with fp32 scales [..., out/128, in/128] -> dtype."""
-    *lead, rows, cols = weight.shape
-    row_blocks, col_blocks = scale_inv.shape[-2:]
-    blocks = weight.view(*lead, row_blocks, FP8_BLOCK_SIZE, col_blocks, FP8_BLOCK_SIZE)
-    scaled = blocks.to(torch.float32) * scale_inv[..., :, None, :, None]
-    return scaled.view(*lead, rows, cols).to(dtype)
-
-
-def clamp_to_float16(tensor: Tensor, name: str) -> Tensor:
-    """Saturate the values float16 cannot hold, logging how many there were."""
-    overflow = int((tensor.abs() > FLOAT16_MAX).sum())
-    if overflow:
-        logger.warning(f"{name}: {overflow} values exceed the float16 range; clamped")
-        tensor = tensor.clamp(-FLOAT16_MAX, FLOAT16_MAX)
-    return tensor
-
-
-def select_fp8(weight: Tensor, index: Tensor | int) -> Tensor:
-    """weight[index] for a float8 tensor; advanced indexing has no float8 kernel."""
-    return weight.view(torch.uint8)[index].view(weight.dtype)
-
-
-class DenseMLP(nn.Module):
-    def __init__(
-        self,
-        hidden: int,
-        intermediate: int,
-        *,
-        dtype: torch.dtype | None = None,
-        device: Device = None,
-    ) -> None:
-        super().__init__()
-        linear = dict(bias=False, dtype=dtype, device=device)
-        self.gate_proj = nn.Linear(hidden, intermediate, **linear)
-        self.up_proj = nn.Linear(hidden, intermediate, **linear)
-        self.down_proj = nn.Linear(intermediate, hidden, **linear)
-
-    def forward(self, x: Tensor) -> Tensor:
-        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
-
-
-class MoEExperts(nn.Module):
-    """Routed experts stacked per layer: w13 [E, 2I, H] (gate then up), w2 [E, H, I].
-
-    With fp8 the weights are float8_e4m3fn buffers with fp32 block scales; on CUDA
-    they run on the W8A8 kernel, elsewhere they are dequantized per forward.
-    """
-
-    def __init__(
-        self,
-        num_experts: int,
-        hidden: int,
-        intermediate: int,
-        *,
-        fp8: bool,
-        dtype: torch.dtype | None = None,
-        device: Device = None,
-    ) -> None:
-        super().__init__()
-        self.num_experts = num_experts
-        self.hidden = hidden
-        self.intermediate = intermediate
-        self.fp8 = fp8
-        w13_shape = (num_experts, 2 * intermediate, hidden)
-        w2_shape = (num_experts, hidden, intermediate)
-        if not fp8:
-            self.w13_weight = nn.Parameter(
-                torch.empty(w13_shape, dtype=dtype, device=device)
-            )
-            self.w2_weight = nn.Parameter(
-                torch.empty(w2_shape, dtype=dtype, device=device)
-            )
-            return
-        if hidden % FP8_BLOCK_SIZE or intermediate % FP8_BLOCK_SIZE:
-            raise ValueError(
-                f"block-FP8 experts need hidden and intermediate sizes divisible by "
-                f"{FP8_BLOCK_SIZE}; got {hidden} and {intermediate}"
-            )
-        for name, (experts, rows, cols) in (("w13", w13_shape), ("w2", w2_shape)):
-            self.register_buffer(
-                f"{name}_weight",
-                torch.empty((experts, rows, cols), dtype=FP8_DTYPE, device=device),
-            )
-            self.register_buffer(
-                f"{name}_scale_inv",
-                torch.empty(
-                    (experts, rows // FP8_BLOCK_SIZE, cols // FP8_BLOCK_SIZE),
-                    dtype=torch.float32,
-                    device=device,
-                ),
-            )
-
-    def expert_weights(
-        self, index: Tensor | int, dtype: torch.dtype
-    ) -> tuple[Tensor, Tensor]:
-        """Gate-up and down weights of the experts index selects, in dtype."""
-        if not self.fp8:
-            return self.w13_weight[index].to(dtype), self.w2_weight[index].to(dtype)
-        return (
-            dequantize_fp8_blocks(
-                select_fp8(self.w13_weight, index), self.w13_scale_inv[index], dtype
-            ),
-            dequantize_fp8_blocks(
-                select_fp8(self.w2_weight, index), self.w2_scale_inv[index], dtype
-            ),
-        )
-
-    def forward(self, x: Tensor, topk_ids: Tensor, topk_weights: Tensor) -> Tensor:
-        """x [N, H], topk_ids / topk_weights [N, k] -> weighted expert sum [N, H]."""
-        if self.fp8 and x.is_cuda:
-            return fp8_moe_forward(
-                x,
-                topk_ids,
-                topk_weights,
-                self.w13_weight,
-                self.w13_scale_inv,
-                self.w2_weight,
-                self.w2_scale_inv,
-            )
-        if topk_ids.numel() <= GATHERED_EXPERT_MAX_ROWS:
-            return self.forward_gathered(x, topk_ids, topk_weights)
-        return self.forward_grouped(x, topk_ids, topk_weights)
-
-    def forward_gathered(
-        self, x: Tensor, topk_ids: Tensor, topk_weights: Tensor
-    ) -> Tensor:
-        """One batched matmul over every (token, expert) pair; no host sync."""
-        tokens, top_k = topk_ids.shape
-        w13, w2 = self.expert_weights(topk_ids.reshape(-1), x.dtype)
-        rows = x.repeat_interleave(top_k, dim=0).unsqueeze(-1)  # [M, H, 1]
-        gate, up = torch.bmm(w13, rows).squeeze(-1).chunk(2, dim=-1)
-        hidden = (F.silu(gate) * up).unsqueeze(-1)
-        out = torch.bmm(w2, hidden).squeeze(-1)  # [M, H]
-        out = out.float() * topk_weights.reshape(-1, 1).float()
-        return out.view(tokens, top_k, self.hidden).sum(dim=1).to(x.dtype)
-
-    def forward_grouped(
-        self, x: Tensor, topk_ids: Tensor, topk_weights: Tensor
-    ) -> Tensor:
-        """Loop over the experts the batch hits, each on its own tokens."""
-        tokens, top_k = topk_ids.shape
-        flat_ids = topk_ids.reshape(-1)
-        order = torch.argsort(flat_ids, stable=True)
-        token_rows = order // top_k
-        pair_weights = topk_weights.reshape(-1)[order]
-        counts = torch.bincount(flat_ids, minlength=self.num_experts).tolist()
-        out = torch.zeros(tokens, self.hidden, dtype=torch.float32, device=x.device)
-        start = 0
-        for expert, count in enumerate(counts):
-            if count == 0:
-                continue
-            rows = token_rows[start : start + count]
-            weights = pair_weights[start : start + count]
-            start += count
-            w13, w2 = self.expert_weights(expert, x.dtype)
-            gate, up = F.linear(x[rows], w13).chunk(2, dim=-1)
-            expert_out = F.linear(F.silu(gate) * up, w2)
-            out.index_add_(0, rows, expert_out.float() * weights[:, None].float())
-        return out.to(x.dtype)
-
-
-class SparseMoeBlock(nn.Module):
-    """Top-k routed experts plus a sigmoid-gated shared expert."""
-
-    def __init__(
-        self,
-        config: FishOmniTextConfig,
-        *,
-        fp8: bool,
-        dtype: torch.dtype | None = None,
-        device: Device = None,
-    ) -> None:
-        super().__init__()
-        hidden = config.hidden_size
-        self.top_k = config.num_experts_per_tok
-        self.norm_topk_prob = config.norm_topk_prob
-        linear = dict(bias=False, dtype=dtype, device=device)
-        self.gate = nn.Linear(hidden, config.num_experts, **linear)
-        self.experts = MoEExperts(
-            config.num_experts,
-            hidden,
-            config.moe_intermediate_size,
-            fp8=fp8,
-            dtype=dtype,
-            device=device,
-        )
-        self.shared_expert = DenseMLP(
-            hidden, config.shared_expert_intermediate_size, dtype=dtype, device=device
-        )
-        self.shared_expert_gate = nn.Linear(hidden, 1, **linear)
-
-    def forward(self, x: Tensor) -> Tensor:
-        flat = x.reshape(-1, x.shape[-1])
-        probabilities = torch.softmax(self.gate(flat).float(), dim=-1)
-        topk_weights, topk_ids = probabilities.topk(self.top_k, dim=-1)
-        if self.norm_topk_prob:
-            topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
-        out = self.experts(flat, topk_ids, topk_weights.to(flat.dtype))
-        shared_gate = torch.sigmoid(self.shared_expert_gate(flat))
-        return (out + shared_gate * self.shared_expert(flat)).view_as(x)

fish_speech/models/text2semantic/fish_omni/tower.py

@@ -1,421 +0,0 @@
-import math
-from collections.abc import Iterable
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from loguru import logger
-from torch import Tensor
-from torch.types import Device
-
-from fish_speech.models.text2semantic.fish_omni.config import (
-    AudioDecoderConfig,
-    FishOmniConfig,
-)
-from fish_speech.models.text2semantic.fish_omni.layers import (
-    ZeroCenteredRMSNorm,
-    grouped_attention,
-)
-
-# Repetition-aware sampling of the main token, fixed by the training recipe: when
-# the first draw repeats one of the last RAS_WINDOW_FRAMES main tokens, a second
-# draw with these parameters is used instead.
-RAS_WINDOW_FRAMES = 16
-RAS_TEMPERATURE = 1.0
-RAS_TOP_K = 30
-RAS_TOP_P = 0.95
-MIN_TEMPERATURE = 1e-5
-TOWER_CHECKPOINT_PREFIXES = {
-    "audio_decoder.codebook_embeddings.": "codebook_embeddings.",
-    "audio_decoder_mtp_shared.layers.": "shared_layers.",
-    "audio_decoder_mtp.layers.": "residual_layers.",
-    "audio_decoder_mtp.heads.": "codebook_heads.",
-}
-
-
-def sample_top_k_top_p(
-    logits: Tensor,
-    temperature: float | Tensor,
-    top_k: int | Tensor,
-    top_p: float | Tensor,
-    *,
-    check_finite: bool = True,
-) -> Tensor:
-    """One id per row; drops the token crossing top_p, except the highest-ranked one.
-
-    Sampling parameters may be 0-dim device tensors (CUDA graphs); check_finite
-    syncs with the host.
-    """
-    if isinstance(temperature, Tensor):
-        temperature = temperature.clamp(min=MIN_TEMPERATURE)
-    else:
-        temperature = max(temperature, MIN_TEMPERATURE)
-    if check_finite and not torch.isfinite(logits).all():
-        raise FloatingPointError("non-finite logits; refusing to sample invalid audio")
-    values, ids = torch.sort(logits.float() / temperature, descending=True, dim=-1)
-    ranks = torch.arange(values.shape[-1], device=values.device)
-    dropped = (values.softmax(-1).cumsum(-1) > top_p) | (ranks >= top_k)
-    dropped[..., 0] = False
-    probabilities = values.masked_fill(dropped, -torch.inf).softmax(-1)
-    noise = torch.empty_like(probabilities).exponential_()
-    choice = (probabilities / noise).argmax(-1, keepdim=True)
-    return ids.gather(-1, choice).squeeze(-1)
-
-
-class TowerAttention(nn.Module):
-    """Gated attention over time; checkpoint wqkv rows are [Q | Gate | K | V]."""
-
-    def __init__(
-        self,
-        config: AudioDecoderConfig,
-        *,
-        dtype: torch.dtype | None = None,
-        device: Device = None,
-    ) -> None:
-        super().__init__()
-        self.config = config
-        q_width = config.n_head * config.head_dim
-        kv_width = config.n_local_heads * config.head_dim
-        linear = dict(bias=False, dtype=dtype, device=device)
-        norm = dict(dtype=dtype, device=device)
-        self.wqkv = nn.Linear(config.dim, 2 * q_width + 2 * kv_width, **linear)
-        self.wo = nn.Linear(q_width, config.dim, **linear)
-        # Always on, whatever audio_decoder_config.attention_qk_norm says: training runs
-        # q/k norm in the audio decoder, but the exporter never writes that field, so
-        # released Drama 3 configs carry its default false. Keep the (1 + weight) scale too:
-        # exported weights are zero-centered, not pre-shifted by one.
-        self.q_norm = ZeroCenteredRMSNorm(config.head_dim, config.norm_eps, **norm)
-        self.k_norm = ZeroCenteredRMSNorm(config.head_dim, config.norm_eps, **norm)
-        self.k_cache: Tensor | None = None
-        self.v_cache: Tensor | None = None
-        # CUDA-graph decode attends this many cache rows; the mask hides the rest.
-        self.static_kv_len: int | None = None
-
-    def setup_cache(self, max_seq_len: int) -> None:
-        shape = (1, self.config.n_local_heads, max_seq_len, self.config.head_dim)
-        self.k_cache = self.wqkv.weight.new_zeros(shape)
-        self.v_cache = self.wqkv.weight.new_zeros(shape)
-
-    def rotate(self, x: Tensor, positions: Tensor) -> Tensor:
-        """Interleaved (non-neox) partial RoPE, computed in fp32."""
-        width = self.config.rotary_dim
-        exponents = torch.arange(0, width, 2, device=x.device, dtype=torch.float32)
-        angles = positions.float()[:, None] * self.config.rope_base ** (
-            -exponents / width
-        )
-        cos, sin = angles.cos()[:, None], angles.sin()[:, None]
-        even, odd = x[..., :width:2].float(), x[..., 1:width:2].float()
-        rotated = torch.stack((even * cos - odd * sin, odd * cos + even * sin), -1)
-        return torch.cat((rotated.flatten(-2).to(x.dtype), x[..., width:]), dim=-1)
-
-    def forward(self, x: Tensor, positions: Tensor) -> Tensor:
-        cfg = self.config
-        length = len(positions)
-        q_width, kv_width = cfg.n_head * cfg.head_dim, cfg.n_local_heads * cfg.head_dim
-        q, gate, k, v = self.wqkv(x).split((q_width, q_width, kv_width, kv_width), -1)
-        q = self.q_norm(q.reshape(length, cfg.n_head, cfg.head_dim))
-        k = self.k_norm(k.reshape(length, cfg.n_local_heads, cfg.head_dim))
-        q = self.rotate(q, positions).transpose(0, 1).unsqueeze(0)
-        k = self.rotate(k, positions).transpose(0, 1).unsqueeze(0)
-        v = v.reshape(length, cfg.n_local_heads, cfg.head_dim)
-        v = v.transpose(0, 1).unsqueeze(0)
-        if self.k_cache is not None:
-            if self.static_kv_len is not None:
-                end = self.static_kv_len
-            else:
-                end = int(positions[-1]) + 1
-                assert end <= self.k_cache.shape[2], "tower cache capacity exceeded"
-            self.k_cache.index_copy_(2, positions, k)
-            self.v_cache.index_copy_(2, positions, v)
-            k, v = self.k_cache[:, :, :end], self.v_cache[:, :, :end]
-            key_positions = torch.arange(end, device=x.device)
-        else:
-            key_positions = positions
-        mask = key_positions[None, :] <= positions[:, None]
-        out = grouped_attention(q, k, v, mask)
-        out = out.squeeze(0).transpose(0, 1).reshape(length, q_width)
-        return self.wo(out * gate.sigmoid())
-
-
-class FP32OutputLinear(nn.Linear):
-    """FP16 GEMM with FP32 output, keeping values beyond the FP16 range."""
-
-    def forward(self, x: Tensor) -> Tensor:
-        if x.dtype != torch.float16:
-            return super().forward(x)
-        if x.is_cuda:
-            out = torch.mm(
-                x.reshape(-1, self.in_features), self.weight.T, out_dtype=torch.float32
-            )
-            out = out.reshape(*x.shape[:-1], self.out_features)
-        else:
-            out = F.linear(x.float(), self.weight.float())
-        return out if self.bias is None else out + self.bias.float()
-
-
-class TowerMLP(nn.Module):
-    def __init__(
-        self,
-        config: AudioDecoderConfig,
-        *,
-        dtype: torch.dtype | None = None,
-        device: Device = None,
-    ) -> None:
-        super().__init__()
-        linear = dict(bias=False, dtype=dtype, device=device)
-        self.w1 = nn.Linear(config.dim, config.intermediate_size, **linear)
-        self.w3 = nn.Linear(config.dim, config.intermediate_size, **linear)
-        # The checkpoint's down projections exceed 65504 under FP16.
-        self.w2 = FP32OutputLinear(config.intermediate_size, config.dim, **linear)
-
-    def forward(self, x: Tensor) -> Tensor:
-        return self.w2(F.silu(self.w1(x)) * self.w3(x))
-
-
-class TowerTransformerLayer(nn.Module):
-    def __init__(
-        self,
-        config: AudioDecoderConfig,
-        *,
-        dtype: torch.dtype | None = None,
-        device: Device = None,
-    ) -> None:
-        super().__init__()
-        args = dict(dtype=dtype, device=device)
-        self.attention = TowerAttention(config, **args)
-        self.attention_norm = ZeroCenteredRMSNorm(config.dim, config.norm_eps, **args)
-        self.ffn_norm = ZeroCenteredRMSNorm(config.dim, config.norm_eps, **args)
-        self.feed_forward = TowerMLP(config, **args)
-
-    def forward(self, x: Tensor, positions: Tensor) -> Tensor:
-        dtype = self.attention.wqkv.weight.dtype
-        residual = x.float() if dtype == torch.float16 else x
-        residual = residual + self.attention(
-            self.attention_norm(residual).to(dtype), positions
-        )
-        return residual + self.feed_forward(self.ffn_norm(residual).to(dtype))
-
-
-class TowerLayer(nn.Module):
-    """Mix a code embedding into the hidden state, then one pre-norm block."""
-
-    def __init__(
-        self,
-        config: AudioDecoderConfig,
-        *,
-        shared: bool,
-        dtype: torch.dtype | None = None,
-        device: Device = None,
-    ) -> None:
-        super().__init__()
-        self.shared = shared
-        args = dict(dtype=dtype, device=device)
-        # Checkpoint name of the embedding differs between shared / residual layers.
-        name = "semantic_embedding" if shared else "codebook_embedding"
-        self.add_module(name, nn.Embedding(config.vocab_size, config.dim, **args))
-        self.enorm = ZeroCenteredRMSNorm(config.dim, config.norm_eps, **args)
-        self.hnorm = ZeroCenteredRMSNorm(config.dim, config.norm_eps, **args)
-        self.eh_proj = nn.Linear(2 * config.dim, config.dim, bias=False, **args)
-        self.transformer_layer = TowerTransformerLayer(config, **args)
-        self.final_layernorm = ZeroCenteredRMSNorm(config.dim, config.norm_eps, **args)
-
-    def forward(
-        self,
-        hidden: Tensor,
-        codes: Tensor,
-        semantic: Tensor,
-        positions: Tensor,
-        shared_input: Tensor | None = None,
-    ) -> Tensor:
-        embedding = self.semantic_embedding if self.shared else self.codebook_embedding
-        codes = codes.clamp(0, embedding.num_embeddings - 1)
-        e = torch.where(semantic[:, None], embedding(codes), 0.0)
-        if shared_input is not None:
-            e = e + shared_input
-        mixed = self.eh_proj(torch.cat((self.enorm(e), self.hnorm(hidden)), dim=-1))
-        out = self.final_layernorm(self.transformer_layer(mixed, positions))
-        return out.to(mixed.dtype)
-
-
-class FishOmniAudioDecoder(nn.Module):
-    """Codebook tower: predicts codes 1..N-1 of each frame, attending over time.
-
-    A frame row is [main token id, code 0, ..., code N-1]; code 0 is the main
-    token's offset in the semantic band.
-    """
-
-    def __init__(
-        self,
-        config: FishOmniConfig,
-        *,
-        dtype: torch.dtype = torch.float16,
-        device: Device = None,
-    ) -> None:
-        super().__init__()
-        cfg = self.config = config.audio_decoder
-        self.semantic_begin = config.semantic_start_token_id
-        self.semantic_end = config.semantic_end_token_id
-        self.eos_id = config.eos_token_id
-        args = dict(dtype=dtype, device=device)
-        self.codebook_embeddings = nn.Embedding(
-            cfg.vocab_size * cfg.num_codebooks, cfg.dim, **args
-        )
-        self.shared_layers = nn.ModuleList(
-            TowerLayer(cfg, shared=True, **args)
-            for _ in range(cfg.shared_enrichment_layers)
-        )
-        self.residual_layers = nn.ModuleList(
-            TowerLayer(cfg, shared=False, **args) for _ in range(cfg.num_codebooks - 1)
-        )
-        self.codebook_heads = nn.ModuleList(
-            nn.Linear(cfg.dim, cfg.vocab_size, bias=False, **args)
-            for _ in range(cfg.num_codebooks - 1)
-        )
-        # Ids the main token is drawn from (eos + the semantic band), built lazily
-        # on the tower's device: creating it inside CUDA-graph capture is illegal.
-        self.main_candidates: Tensor | None = None
-
-    @property
-    def device(self) -> torch.device:
-        return self.codebook_embeddings.weight.device
-
-    def setup_caches(self, max_seq_len: int) -> None:
-        for layer in (*self.shared_layers, *self.residual_layers):
-            layer.transformer_layer.attention.setup_cache(max_seq_len)
-
-    def is_semantic(self, tokens: Tensor) -> Tensor:
-        return (tokens >= self.semantic_begin) & (tokens <= self.semantic_end)
-
-    def embedding_sum(self, frames: Tensor, first: int, end: int) -> Tensor:
-        """Summed embeddings of codes first..end-1 per frame."""
-        if first == end:
-            return self.codebook_embeddings.weight.new_zeros(
-                (len(frames), self.config.dim)
-            )
-        offsets = torch.arange(first, end, device=self.device) * self.config.vocab_size
-        rows = frames[:, 1 + first : 1 + end] + offsets
-        return F.embedding_bag(rows, self.codebook_embeddings.weight, mode="sum")
-
-    def embed_backbone(self, frames: Tensor, text_embeddings: Tensor) -> Tensor:
-        """Backbone input: semantic frames add codes 1..K to the token embedding."""
-        frames = frames.to(self.device)
-        count = self.config.vq_backbone_codebooks
-        extra = self.embedding_sum(frames, 1, 1 + count).to(text_embeddings.device)
-        mixed = (text_embeddings + extra) / math.sqrt(1 + count)
-        semantic = self.is_semantic(frames[:, 0]).to(text_embeddings.device)
-        return torch.where(semantic[:, None], mixed, text_embeddings)
-
-    def run_shared(
-        self,
-        hidden: Tensor,
-        c0: Tensor,
-        semantic: Tensor,
-        positions: Tensor,
-        frames: Tensor,
-    ) -> Tensor:
-        # The first shared layer also sees the codes the backbone did not.
-        unseen = self.embedding_sum(
-            frames, 1 + self.config.vq_backbone_codebooks, self.config.num_codebooks
-        )
-        unseen = torch.where(self.is_semantic(frames[:, 0])[:, None], unseen, 0.0)
-        for index, layer in enumerate(self.shared_layers):
-            hidden = layer(
-                hidden, c0, semantic, positions, unseen if index == 0 else None
-            )
-        return hidden
-
-    def prefill(
-        self, hidden: Tensor, positions: Tensor, frames: Tensor, next_frames: Tensor
-    ) -> None:
-        """Teacher-force the tower cache; row p is labelled with frame p + 1."""
-        semantic = self.is_semantic(next_frames[:, 0])
-        hidden = self.run_shared(hidden, next_frames[:, 1], semantic, positions, frames)
-        for index, layer in enumerate(self.residual_layers):
-            hidden = layer(hidden, next_frames[:, index + 1], semantic, positions)
-
-    def decode_step(
-        self,
-        logits: Tensor,
-        hidden: Tensor,
-        positions: Tensor,
-        frames: Tensor,
-        history: Tensor,
-        *,
-        temperature: float | Tensor,
-        top_k: int | Tensor,
-        top_p: float | Tensor,
-        check_finite: bool = True,
-    ) -> Tensor:
-        """Sample the frame after each row; the main token is drawn over eos + semantics."""
-        if self.main_candidates is None or self.main_candidates.device != self.device:
-            eos = torch.tensor([self.eos_id], device=self.device)
-            band = torch.arange(
-                self.semantic_begin, self.semantic_end + 1, device=self.device
-            )
-            self.main_candidates = torch.cat((eos, band)).sort().values
-        candidates = self.main_candidates
-        candidate_logits = logits[:, candidates]
-        first = candidates[
-            sample_top_k_top_p(
-                candidate_logits, temperature, top_k, top_p, check_finite=check_finite
-            )
-        ]
-        retry = candidates[
-            sample_top_k_top_p(
-                candidate_logits,
-                RAS_TEMPERATURE,
-                RAS_TOP_K,
-                RAS_TOP_P,
-                check_finite=False,
-            )
-        ]
-        # Only audio tokens are resampled, as in S2: the window starts with prompt rows,
-        # which hold the user turn's eos, and a stop must not be redrawn.
-        repeated = (history[-RAS_WINDOW_FRAMES:] == first[:, None]).any(-1)
-        repeated &= self.is_semantic(first)
-        main = torch.where(repeated, retry, first)
-        c0 = (main - self.semantic_begin).clamp(0, self.config.vocab_size - 1)
-        semantic = self.is_semantic(main)
-        hidden = self.run_shared(hidden, c0, semantic, positions, frames)
-        codes = [c0]
-        for layer, head in zip(self.residual_layers, self.codebook_heads):
-            hidden = layer(hidden, codes[-1], semantic, positions)
-            codes.append(
-                sample_top_k_top_p(
-                    head(hidden), temperature, top_k, top_p, check_finite=check_finite
-                )
-            )
-        return torch.stack((main, *codes), dim=-1)
-
-
-@torch.no_grad()
-def load_audio_decoder_weights(
-    model: FishOmniAudioDecoder, weights: Iterable[tuple[str, Tensor]]
-) -> None:
-    """Copy audio_decoder* checkpoint tensors into model, each parameter exactly once."""
-    targets = dict(model.named_parameters())
-    loaded = set()
-    for key, tensor in weights:
-        name = next(
-            (
-                new + key[len(old) :]
-                for old, new in TOWER_CHECKPOINT_PREFIXES.items()
-                if key.startswith(old)
-            ),
-            None,
-        )
-        if name not in targets:
-            raise ValueError(f"unknown tower weight {key}")
-        if name in loaded:
-            raise ValueError(f"tower weight filled twice: {key}")
-        target = targets[name]
-        if target.shape != tensor.shape:
-            raise ValueError(f"{key}: shape {tensor.shape} != {target.shape}")
-        value = tensor.to(device=target.device, dtype=target.dtype)
-        if not torch.isfinite(value).all():
-            raise FloatingPointError(f"{key} cannot be represented in {target.dtype}")
-        target.copy_(value)
-        loaded.add(name)
-    if missing := targets.keys() - loaded:
-        raise ValueError(f"tower weights not loaded: {sorted(missing)[:8]}")
-    logger.info(f"FishOmni tower: loaded {len(loaded)} tensors on {model.device}")

fish_speech/models/text2semantic/sampling.py

@@ -1,9 +0,0 @@
-"""Sampling defaults for S2-Pro and Drama 3: those of the official Fish Audio SDK
-(fish-audio-sdk TTSConfig). The SDK has no top_k; both models use 30."""
-
-DEFAULT_MAX_NEW_TOKENS = 1024  # per text chunk
-DEFAULT_TEMPERATURE = 0.7
-DEFAULT_TOP_P = 0.7
-DEFAULT_TOP_K = 30
-DEFAULT_REPETITION_PENALTY = 1.2
-DEFAULT_CHUNK_LENGTH = 200

footer.html

@@ -0,0 +1,11 @@
+<section class="card">
+
+    <section class="credits" aria-labelledby="acknowledgments-title">
+      <h2 id="acknowledgments-title">Acknowledgments</h2>
+      <p>It took us over half a year to bring SGLang Omni back to life. Thanks to all the people who helped us along the way.</p>
+      <p>Thanks so much to <a href="https://huggingface.co/andito" target="_blank" rel="noreferrer">andito</a>, <a href="https://x.com/_akhaliq" target="_blank" rel="noreferrer">AK</a>, and <a href="https://huggingface.co/multimodalart" target="_blank" rel="noreferrer">poli</a> for helping set up the Space.</p>
+      <p><b>SGLang Omni team:</b> <a href="https://x.com/GenAI_is_real" target="_blank" rel="noreferrer">Chenyang Zhao</a>, Jingwen Gu, Shenggui Li, Richard Wang, Yichi Zhang, Xuesong Ye, Jiaxin Deng, Ratish P, Jiaxuan Luo, Xinyu Lu, Wenyao Gao, Ao Sun, Jintao Quaige Li, Yuhao Chen, Yuhan Zhou, Xuxiang Tian, Ruoyi Zhang, Junnan Li, Yixiao Zeng, Chenchen Hong, Huapeng Zhou, Yang Gao, Chaoye Fei, Yiyang Zh, Kang, Yitian Gong, Yuan Luo, Hao Jin, Yifei Gao, Gaokai Zhang, Lifan Shen, Hui Shen, Kejing Xia, Jinxin Li, Dongming Shen, Jun Liu, Yue Yin, Zhihao Guo, Jisen Li, Ruoyu He, Xuanyi Li, Chenye Zhu, Liang Zhen, Xiaodong Ye, Yizhi Zhen, Kaizuo Li, Chonghan Liu, Wei Guan, Zhuofeng Li</p>
+      <p><b>FishAudio Team:</b> Shijia Liao, Shidong Li, Xuehao Yang, Jiahua Liu</p>
+    </section>
+    <p class="fine" style="margin-top:10px">Model license: <a href="https://huggingface.co/fishaudio/drama-3/blob/main/LICENSE" target="_blank" rel="noreferrer">Fish Audio Community License</a>, for research and non-commercial use. Takes and uploaded voices live in server memory only and expire.</p>
+  </section>

gradio_ui.py

@@ -0,0 +1,479 @@
+"""Gradio form, classic streaming stage, and a separately mounted gr.Workflow canvas."""
+import asyncio
+import base64
+from contextlib import aclosing
+import html
+import json
+import os
+import secrets
+from pathlib import Path
+
+import anyio
+from fastapi import HTTPException
+import gradio as gr
+import numpy as np
+from pydantic import ValidationError
+from starlette.middleware.sessions import SessionMiddleware
+from starlette.datastructures import UploadFile
+
+ROOT = Path(__file__).resolve().parent
+OUTPUT_DIR = Path(os.environ.get("GRADIO_TEMP_DIR", "/tmp/drama3-gradio"))
+PRESETS = json.loads((ROOT / "voices/voices.json").read_text())
+CONTENT = json.loads((ROOT / "studio_content.json").read_text())
+DEFAULT_DESIGN = CONTENT["CAST"][0]
+
+
+async def events(response):
+    """Consume SSE without running a second HTTP server or bypassing admission."""
+    pending = ""
+    async with aclosing(response.body_iterator):
+        async for chunk in response.body_iterator:
+            pending += chunk.decode() if isinstance(chunk, bytes) else chunk
+            while "\n\n" in pending:
+                raw, pending = pending.split("\n\n", 1)
+                for line in raw.splitlines():
+                    if line.startswith("data: "):
+                        yield json.loads(line[6:])
+
+
+class Studio:
+    def __init__(self, backend):
+        self.backend = backend
+
+    async def identity(self, request):
+        if request is None or request.request is None:
+            raise gr.Error("Open the studio in your browser before generating a take.")
+        try:
+            return await asyncio.to_thread(self.backend.identify, request.request)
+        except HTTPException as error:
+            raise gr.Error(str(error.detail)) from error
+
+    async def response(self, body, request):
+        try:
+            identity = await self.identity(request)
+            if os.environ.get("DEMO_PREVIEW") == "1":
+                raise gr.Error("This is a UI preview. Start the GPU engine to perform a take.")
+            return await self.backend.generate(body, identity)
+        except HTTPException as error:
+            raise gr.Error(str(error.detail)) from error
+
+    def body(self, text, mode, voice, transcript, name, description, scene,
+             temperature, top_p, top_k, retake_of=None, note=None):
+        try:
+            return self.backend.GenerateBody(
+                text=text, voice=self.backend.VoiceBody(
+                    mode={"Clone": "clone", "Design": "design", "Surprise": "random"}[mode],
+                    id=voice, transcript=transcript, name=name,
+                    description=description, scene=scene),
+                temperature=temperature, top_p=top_p, top_k=top_k,
+                retake_of=retake_of, note=note)
+        except (ValidationError, KeyError) as error:
+            raise gr.Error("Check the voice mode and sampling settings.") from error
+
+    async def run(self, body, current, history, request):
+        response = await self.response(body, request)
+        rate = self.backend.SAMPLE_RATE
+        pending = bytearray()
+        done = False
+        async with aclosing(events(response)) as stream:
+            async for event in stream:
+                kind = event["type"]
+                if kind == "error":
+                    raise gr.Error(event["message"])
+                if kind == "queued":
+                    yield gr.skip(), gr.skip(), gr.skip(), f"Waiting for the stage · position {event['position']}", gr.skip()
+                elif kind == "retrying":
+                    yield gr.skip(), gr.skip(), gr.skip(), event["message"], gr.skip()
+                elif kind == "start":
+                    rate = event["sample_rate"]
+                    yield None, gr.skip(), gr.skip(), "Preparing the actor…", gr.skip()
+                elif kind == "audio":
+                    pending.extend(base64.b64decode(event["pcm"]))
+                    # Gradio re-encodes each chunk for its native streaming player.
+                    # Half-second chunks avoid a re-encode per 46 ms engine frame.
+                    if len(pending) >= rate:
+                        chunk = self.backend.wav_bytes(self.backend.pcm_to_float(bytes(pending)), rate)
+                        pending.clear()
+                        yield chunk, gr.skip(), gr.skip(), "On air…", gr.skip()
+                elif kind == "done":
+                    done = True
+                    audio = (self.backend.wav_bytes(self.backend.pcm_to_float(bytes(pending)), rate)
+                             if pending else gr.skip())
+                    entries = list(history or []) + [event["take_id"]]
+                    choices = [(f"Take {i + 1}", key) for i, key in enumerate(entries)]
+                    metrics = (f"First audio **{event['ttfa_ms']} ms** · "
+                               f"**{event['audio_s']:.1f} s** audio · **{event['speed']}×** real time")
+                    if event.get("warning"):
+                        metrics += "\n\n" + event["warning"]
+                    yield audio, event["take_id"], entries, metrics, gr.Dropdown(choices=choices, value=event["take_id"])
+        if not done:
+            raise gr.Error("The take ended before it finished. Try again.")
+
+    async def stage_run(self, body, current, history, request):
+        """Stream a PCM packet ledger to the custom Gradio stage.
+
+        Each update includes all packets so browser render batching cannot lose a
+        chunk. Gradio sends array diffs, and the player consumes each packet once.
+        """
+        previous = self.backend.store.take(body.retake_of) if body.retake_of else None
+        choice = previous.choice if previous else None
+        mode = choice.mode if choice else body.voice.mode
+        voice = choice.voice if choice and mode == "clone" else self.backend.store.voice(body.voice.id or "")
+        character = choice.character.name if choice and choice.character else body.voice.name
+        label = voice.label if mode == "clone" and voice else character if mode == "design" else "Surprise"
+        common = {"run_id": secrets.token_urlsafe(9), "text": body.text, "mode": mode,
+                  "voice": label or "Narrator", "note": body.note if previous else None}
+        packets, pending = [], bytearray()
+        rate, sequence, done = self.backend.SAMPLE_RATE, 0, False
+
+        def update(event):
+            nonlocal sequence
+            sequence += 1
+            return {**common, **event, "sequence": sequence, "packets": list(packets)}
+
+        yield update({"type": "preparing"}), gr.skip(), gr.skip()
+        try:
+            response = await self.response(body, request)
+            async with aclosing(events(response)) as stream:
+                async for event in stream:
+                    kind = event["type"]
+                    if kind == "start":
+                        rate = event["sample_rate"]
+                    if kind == "audio":
+                        pending.extend(base64.b64decode(event["pcm"]))
+                        if len(pending) < rate:
+                            continue
+                        packets.append(base64.b64encode(pending).decode("ascii"))
+                        pending.clear()
+                        event = {k: v for k, v in event.items() if k != "pcm"}
+                        event["sample_rate"] = rate
+                    elif kind in {"done", "error"}:
+                        if pending:
+                            packets.append(base64.b64encode(pending).decode("ascii"))
+                            pending.clear()
+                        event["sample_rate"] = rate
+                        done = True
+                    if kind == "done":
+                        entries = list(history or []) + [event["take_id"]]
+                        yield update(event), event["take_id"], entries
+                    else:
+                        yield update(event), gr.skip(), gr.skip()
+            if not done:
+                yield update({"type": "error", "message": "The take ended before it finished. Try again."}), gr.skip(), gr.skip()
+        except gr.Error as error:
+            yield update({"type": "error", "message": str(error)}), gr.skip(), gr.skip()
+
+    async def stage_perform(self, text, mode, voice, transcript, name, description, scene,
+                            temperature, top_p, top_k, current, history, request: gr.Request):
+        try:
+            body = self.body(text, mode, voice, transcript, name, description, scene, temperature, top_p, top_k)
+        except gr.Error as error:
+            yield {"type": "error", "message": str(error)}, gr.skip(), gr.skip()
+            return
+        async with aclosing(self.stage_run(body, current, history, request)) as stream:
+            async for result in stream:
+                yield result
+
+    async def stage_retake(self, text, mode, voice, transcript, name, description, scene,
+                           temperature, top_p, top_k, current, history, note, request: gr.Request):
+        try:
+            if not current:
+                raise gr.Error("Perform a take first, then give the actor a note.")
+            body = self.body(text, mode, voice, transcript, name, description, scene,
+                             temperature, top_p, top_k, current, note)
+        except gr.Error as error:
+            yield {"type": "error", "message": str(error)}, gr.skip(), gr.skip()
+            return
+        async with aclosing(self.stage_run(body, current, history, request)) as stream:
+            async for result in stream:
+                yield result
+
+    async def stage_action(self, choices, evt: gr.EventData, request: gr.Request):
+        """Native Gradio events from the stage's keep/replay controls."""
+        try:
+            await self.identity(request)
+            take_id = evt.take_id
+            if evt.action == "keep":
+                kept = await self.keep(take_id, choices, request)
+                return (*kept, gr.skip(), gr.skip(), {"type": "notice", "message": "Saved as a clone voice. New lines will keep this voice."})
+            if evt.action == "select":
+                take = self.backend.store.take(take_id)
+                if take is None:
+                    raise gr.Error("That take has expired. Perform a fresh one.")
+                return (*[gr.skip() for _ in range(8)], take_id, take.line, {"type": "selected", "take_id": take_id})
+            raise gr.Error("Unknown stage action.")
+        except gr.Error as error:
+            return (*[gr.skip() for _ in range(10)], {"type": "notice", "error": True, "message": str(error)})
+
+    async def perform(self, text, mode, voice, transcript, name, description, scene,
+                      temperature, top_p, top_k, current, history, request: gr.Request):
+        body = self.body(text, mode, voice, transcript, name, description, scene, temperature, top_p, top_k)
+        async with aclosing(self.run(body, current, history, request)) as stream:
+            async for result in stream:
+                yield result
+
+    async def retake(self, text, mode, voice, transcript, name, description, scene,
+                     temperature, top_p, top_k, current, history, note, request: gr.Request):
+        if not current:
+            raise gr.Error("Perform a take first, then give the actor a note.")
+        body = self.body(text, mode, voice, transcript, name, description, scene,
+                         temperature, top_p, top_k, current, note)
+        async with aclosing(self.run(body, current, history, request)) as stream:
+            async for result in stream:
+                yield result
+
+    async def voice_details(self, voice_id, request: gr.Request):
+        await self.identity(request)
+        voice = self.backend.store.voice(voice_id or "")
+        if voice is None:
+            raise gr.Error("That voice has expired. Upload it again.")
+        wav, rate = self.backend.decode_audio(voice.wav)
+        return (rate, wav), voice.transcript
+
+    async def upload(self, path, choices, request: gr.Request):
+        await self.identity(request)
+        if not path:
+            return gr.skip(), gr.skip(), gr.skip(), gr.skip()
+        path = Path(path)
+        if path.stat().st_size > self.backend.UPLOAD_MAX_BYTES:
+            raise gr.Error("Use an audio clip under 25 MB.")
+        try:
+            with path.open("rb") as handle:
+                voice = await self.backend.add_voice(UploadFile(handle, filename=path.name))
+        except HTTPException as error:
+            raise gr.Error(str(error.detail)) from error
+        return self.voice_update(voice, choices)
+
+    def voice_update(self, voice, choices):
+        entries = list(choices or [])
+        pair = (voice["label"], voice["id"])
+        entries = [p for p in entries if p[1] != voice["id"]] + [pair]
+        stored = self.backend.store.voice(voice["id"])
+        wav, rate = self.backend.decode_audio(stored.wav)
+        return gr.Radio(choices=entries, value=voice["id"]), (rate, wav), voice["transcript"], entries
+
+    async def keep(self, take_id, choices, request: gr.Request):
+        await self.identity(request)
+        if not take_id:
+            raise gr.Error("Perform a take first, then keep its voice.")
+        try:
+            voice = await self.backend.voice_from_take(self.backend.FromTake(take_id=take_id))
+        except HTTPException as error:
+            raise gr.Error(str(error.detail)) from error
+        return (*self.voice_update(voice, choices), gr.Radio(value="Clone"), gr.Column(visible=True),
+                gr.Column(visible=False), gr.Column(visible=False))
+
+    def take_path(self, take):
+        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+        path = OUTPUT_DIR / f"drama3-{take.id}.wav"
+        if not path.exists():
+            path.write_bytes(self.backend.wav_bytes(self.backend.pcm_to_float(take.pcm), take.sample_rate))
+        return str(path)
+
+    async def recall(self, take_id, request: gr.Request):
+        await self.identity(request)
+        take = self.backend.store.take(take_id or "")
+        if take is None:
+            raise gr.Error("That take has expired. Perform a fresh one.")
+        return self.take_path(take), take_id, take.line
+
+    async def status(self, request: gr.Request):
+        try:
+            identity = await self.identity(request)
+            state = await self.backend.status(identity)
+        except gr.Error as error:
+            return (f'<div class="hero-top"><span class="pill">{html.escape(str(error))}</span>'
+                    '<a class="pill signin" href="/oauth/huggingface/login" target="_blank" rel="noopener">Sign in with Hugging Face</a></div>',
+                    gr.Button(interactive=False), gr.skip())
+        engine = state["engine"]
+        ready = engine["state"] == "ready"
+        phase = "UI preview · GPU engine offline" if os.environ.get("DEMO_PREVIEW") == "1" else engine["phase"]
+        gpu = state.get("gpu") or {}
+        label = gpu.get("name") or "Drama 3"
+        usage = state["usage"]
+        quota = "Unlimited takes" if usage["unlimited"] else f"{usage['remaining']} takes left today"
+        if identity.signed_in:
+            account = (f'<details class="account"><summary class="pill">{html.escape(identity.name)} · {quota}</summary>'
+                       '<a href="/oauth/huggingface/logout">Sign out</a></details>')
+        elif os.environ.get("SPACE_ID"):
+            account = f'<a class="pill signin" href="/oauth/huggingface/login" target="_blank" rel="noopener">Sign in · {quota}</a>'
+        else:
+            account = '<span class="pill">Local studio</span>'
+        p50 = state["stats"]["ttfa_ms_p50"]
+        timing = f" · {p50} ms first audio" if p50 is not None else ""
+        banner = (f'<div class="hero-top"><span class="pill"><i class="dot {"ready" if ready else "busy"}"></i>'
+                  f'{html.escape(label)}{timing}</span><span class="spacer"></span>{account}</div>')
+        boot = gr.HTML(value=f'<div class="boot"><b>Setting the stage…</b><p>{html.escape(phase)}</p></div>', visible=not ready)
+        return banner, gr.Button(value="● Perform" if ready else "Warming up…", interactive=ready), boot
+
+    async def workflow_generate(self, text, mode, voice_id, description, request):
+        body = self.body(text, mode, voice_id, None, "Narrator", description, None, 1.0, 0.9, 0)
+        response = await self.response(body, request)
+        take = None
+        async with aclosing(events(response)) as stream:
+            async for event in stream:
+                if event["type"] == "error":
+                    raise gr.Error(event["message"])
+                if event["type"] == "done":
+                    take = self.backend.store.take(event["take_id"])
+        if take is None:
+            raise gr.Error("No completed take came back from the engine.")
+        path = self.take_path(take)
+        from gradio_client.utils import encode_file_path
+        return {"path": path, "url": f"/workflow/gradio_api/file={encode_file_path(path)}", "is_file": True}
+
+
+def build_apps(backend):
+    studio = Studio(backend)
+    voice_choices = [(entry["label"], entry["id"]) for entry in PRESETS]
+    with gr.Blocks(title="Fish Audio Drama 3", analytics_enabled=False, fill_width=True, delete_cache=(3600, 86400)) as demo:
+        current = gr.State(None)
+        history = gr.State([])
+        voices = gr.State(voice_choices)
+        with gr.Column(elem_id="studio"):
+            with gr.Row(elem_id="topbar"):
+                banner = gr.HTML(scale=12)
+                theme = gr.Button("◐", min_width=32, scale=0, elem_id="theme-toggle", size="sm")
+            gr.HTML((ROOT / "hero.html").read_text())
+            with gr.Tabs(elem_id="app-tabs"):
+                with gr.Tab("Studio"):
+                    boot = gr.HTML()
+                    with gr.Column(elem_classes="card"):
+                        gr.HTML('<div class="card-head"><h2>Script</h2></div>')
+                        with gr.Row(elem_classes="examples"):
+                            examples = [gr.Button(item["label"], size="sm", min_width=85) for item in CONTENT["EXAMPLES"]]
+                        script = gr.Textbox(value=CONTENT["EXAMPLES"][0]["text"], label="Script", show_label=False,
+                                            lines=4, max_lines=10, elem_id="script", placeholder="Write a line. Give it some direction.")
+                        with gr.Accordion("Direction tags", open=False, elem_id="direction-tags"):
+                            count = gr.Markdown("Tags in `[ ]` and `< >` are direction, not words. Any wording works.", elem_classes="hint")
+                            palette = []
+                            for title, key, cls in [("Events", "POINT_TAGS", "point"), ("Span", "SCOPED_TAGS", "scope"), ("Mood", "STYLES", "mood")]:
+                                with gr.Row(elem_classes=["palette", cls]):
+                                    gr.HTML(f'<span class="plabel">{title}</span>', min_width=60, scale=0)
+                                    for tag in CONTENT[key]:
+                                        button = gr.Button(tag, size="sm", min_width=40, elem_classes="chip")
+                                        palette.append((button, tag, cls))
+                        with gr.Accordion("How to direct", open=False):
+                            gr.Markdown('**Point tag:** `Well [sighs] fine, you win. [laughs]`\n\n'
+                                        '**Scoped tag:** `Promise me <whisper>you won’t tell her</whisper>.`\n\n'
+                                        '**Style prefix:** `[icy calm] Fine.`\n\n'
+                                        '**Director’s note:** `Direction: put the rise on "goodbye".` followed by `Line: So this is goodbye.`')
+                    with gr.Column(elem_classes="card"):
+                        gr.HTML('<div class="card-head"><h2>Voice</h2></div>')
+                        mode = gr.Radio([("🎲 Surprise", "Surprise"), ("🎙 Clone", "Clone"), ("✨ Design", "Design")],
+                                        value="Surprise", show_label=False, container=False, elem_id="voice-mode")
+                        with gr.Column(visible=False) as clone_panel:
+                            voice = gr.Radio(voice_choices, value="clone1", show_label=False, container=False, elem_id="voice-presets")
+                            preview = gr.Audio(label="Reference voice", interactive=False, type="numpy", buttons=[])
+                            with gr.Accordion("Upload or record a voice", open=False):
+                                upload = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Upload or record a voice", format="wav")
+                            transcript = gr.Textbox(label="What the clip says", lines=2, value=PRESETS[0]["transcript"])
+                            gr.Markdown("Clips are trimmed to 20 s and transcribed for you. Only clone voices you have permission to use.", elem_classes="hint")
+                        with gr.Column(visible=False) as design_panel:
+                            with gr.Row(elem_classes="examples"):
+                                cast_buttons = [gr.Button(item["label"], size="sm", min_width=90) for item in CONTENT["CAST"]]
+                            with gr.Row():
+                                name = gr.Textbox(label="Character", value=DEFAULT_DESIGN["name"], scale=1)
+                                description = gr.Textbox(label="Voice description", lines=2, value=DEFAULT_DESIGN["desc"], scale=3)
+                            scene = gr.Textbox(label="Scene (optional)", value=DEFAULT_DESIGN["scene"])
+                            gr.Markdown("Describe age, timbre, accent, and energy. Like a take? **Keep this voice** turns it into a clone.", elem_classes="hint")
+                        with gr.Column() as surprise_panel:
+                            gr.Markdown("Drama 3 improvises a voice that fits the line. Every take is a stranger, so hit **Keep this voice** when you meet one you like.", elem_classes="hint")
+                    with gr.Accordion("Sampling", open=False):
+                        temperature = gr.Slider(.3, 1.3, value=1.0, step=.05, label="Temperature")
+                        top_p = gr.Slider(.5, 1, value=.9, step=.01, label="Top-p")
+                        top_k = gr.Slider(0, 100, value=0, step=1, label="Top-k",
+                                          info="0 samples from every code. Low values make takes go silent mid-line.")
+                        reset = gr.Button("Reset sampling", size="sm")
+                    with gr.Row():
+                        perform = gr.Button("Warming up…", variant="primary", interactive=False, elem_id="perform", scale=5)
+                        stop = gr.Button("Stop", scale=1, elem_id="stop")
+                    with gr.HTML(value={"type": "idle", "packets": []},
+                                 html_template=(ROOT / "stage.html").read_text(),
+                                 css_template=(ROOT / "stage.css").read_text(),
+                                 js_on_load=(ROOT / "stage.js").read_text(),
+                                 apply_default_css=False, elem_id="the-stage") as stage:
+                        with gr.Column(elem_classes="stage-direct"):
+                            gr.HTML('<h3>🎬 Give a note</h3><p class="hint">The actor hears this take, then redoes the line with your note.</p>')
+                            with gr.Row(elem_classes=["examples", "note-chips"]):
+                                note_buttons = [gr.Button(note, size="sm", min_width=60) for note in CONTENT["NOTES"]]
+                            with gr.Row(elem_classes="stage-note-row"):
+                                note = gr.Textbox(show_label=False, placeholder="e.g. more amused than before; drop the hurt", scale=4, elem_id="stage-note")
+                                retake = gr.Button("Retake", variant="primary", scale=1, elem_id="retake")
+                    # Keep the existing native audio API for clients and Workflow.
+                    audio = gr.Audio(visible=False, streaming=True, format="wav", interactive=False)
+                    metrics = gr.Markdown(visible=False)
+                    slate = gr.Dropdown(choices=[], visible=False)
+                    api_perform = gr.Button(visible=False)
+                    api_retake = gr.Button(visible=False)
+                    api_keep = gr.Button(visible=False)
+                with gr.Tab("Workflow", id="workflow"):
+                    gr.HTML('<iframe title="Drama 3 Workflow" src="/workflow/" loading="lazy" allow="microphone; autoplay; clipboard-write" style="width:100%;height:780px;border:0;border-radius:16px"></iframe>')
+                    gr.Markdown('[Open workflow full screen](/workflow/)', elem_classes="hint")
+            gr.HTML((ROOT / "footer.html").read_text())
+
+        for button, item in zip(examples, CONTENT["EXAMPLES"]):
+            button.click(lambda text=item["text"]: text, outputs=script, queue=False, api_visibility="private")
+        for button, item in zip(cast_buttons, CONTENT["CAST"]):
+            button.click(lambda item=item: (item["name"], item["desc"], item["scene"], item["text"]),
+                         outputs=[name, description, scene, script], queue=False, api_visibility="private")
+        for button, value in zip(note_buttons, CONTENT["NOTES"]):
+            button.click(lambda value=value: value, outputs=note, queue=False, api_visibility="private")
+        for button, tag, kind in palette:
+            # Insert at the native textbox's selection and dispatch an input event.
+            button.click(None, inputs=script, outputs=script, queue=False,
+                         js=f"(text) => window.dramaInsert(text, {json.dumps(tag)}, {json.dumps(kind)})", api_visibility="private")
+        script.input(lambda text: f"Tags in `[ ]` and `< >` are direction, not words. · **{len(text)} / 600**",
+                     inputs=script, outputs=count, queue=False, api_visibility="private")
+        mode.change(lambda m: tuple(gr.Column(visible=m == value) for value in ["Clone", "Design", "Surprise"]),
+                    inputs=mode, outputs=[clone_panel, design_panel, surprise_panel], queue=False, api_visibility="private")
+        voice.change(studio.voice_details, inputs=voice, outputs=[preview, transcript], api_visibility="private")
+        upload.change(studio.upload, inputs=[upload, voices], outputs=[voice, preview, transcript, voices], api_visibility="private")
+        stage.click(studio.stage_action, inputs=voices,
+                    outputs=[voice, preview, transcript, voices, mode, clone_panel, design_panel, surprise_panel, current, script, stage],
+                    queue=False, api_visibility="private", show_progress="hidden")
+        slate.input(studio.recall, inputs=slate, outputs=[audio, current, script], api_visibility="private")
+        reset.click(lambda: (1.0, .9, 0), outputs=[temperature, top_p, top_k], queue=False, api_visibility="private")
+        theme.click(None, js="() => window.dramaToggleTheme()", queue=False, api_visibility="private")
+        inputs = [script, mode, voice, transcript, name, description, scene, temperature, top_p, top_k, current, history]
+        outputs = [audio, current, history, metrics, slate]
+        unlock = "(...args) => { window.dramaStage?.unlock(); return args; }"
+        generation = perform.click(studio.stage_perform, inputs=inputs, outputs=[stage, current, history],
+                                   js=unlock, concurrency_limit=None, api_visibility="private", show_progress="hidden")
+        directing = retake.click(studio.stage_retake, inputs=[*inputs, note], outputs=[stage, current, history],
+                                  js=unlock, concurrency_limit=None, api_visibility="private", show_progress="hidden")
+        api_perform.click(studio.perform, inputs=inputs, outputs=outputs, concurrency_limit=None, api_name="perform", show_progress="hidden")
+        api_retake.click(studio.retake, inputs=[*inputs, note], outputs=outputs, concurrency_limit=None, api_name="retake", show_progress="hidden")
+        api_keep.click(studio.keep, inputs=[current, voices], outputs=[voice, preview, transcript, voices, mode, clone_panel, design_panel, surprise_panel], api_name="keep_voice", show_progress="hidden")
+        stop.click(None, cancels=[generation, directing], queue=False, api_visibility="private", js="() => window.dramaStage?.stop()")
+        demo.load(studio.voice_details, inputs=voice, outputs=[preview, transcript], api_visibility="private")
+        demo.load(studio.status, outputs=[banner, perform, boot], api_visibility="private", show_progress="hidden")
+        gr.Timer(5).tick(studio.status, outputs=[banner, perform, boot], queue=False, api_visibility="private", show_progress="hidden")
+
+    # Workflow invokes functions from AnyIO workers. Bridge back to the app loop:
+    # the existing HTTP connection pool and admission events belong to that loop.
+    def direct_performance(text: str, mode: str, voice_id: str, description: str, request: gr.Request) -> dict:
+        return anyio.from_thread.run(studio.workflow_generate, text, mode, voice_id, description, request)
+
+    graph_path = Path(os.environ.get("WORKFLOW_GRAPH", str(OUTPUT_DIR / "workflow.json")))
+    graph_path.parent.mkdir(parents=True, exist_ok=True)
+    if not graph_path.exists():
+        graph_path.write_text((ROOT / "workflow.json").read_text())
+    workflow = gr.Workflow(graph=str(graph_path), bind={"direct_performance": direct_performance})
+    demo.queue(default_concurrency_limit=backend.MAX_ACTIVE, max_size=backend.MAX_QUEUE)
+    workflow.queue(default_concurrency_limit=backend.MAX_ACTIVE, max_size=backend.MAX_QUEUE)
+    return demo, workflow
+
+
+def mount_apps(backend):
+    demo, workflow = build_apps(backend)
+    app = gr.mount_gradio_app(backend.app, workflow, path="/workflow", allowed_paths=[str(OUTPUT_DIR)],
+                              footer_links=[], max_file_size="25mb", ssr_mode=False)
+    # The parent already signs the HF session. A second signer would hide the
+    # logged-in account from quotas and break the shared Studio/Workflow login.
+    workflow.app.user_middleware = [m for m in workflow.app.user_middleware if m.cls is not SessionMiddleware]
+    workflow.app.middleware_stack = None
+    app = gr.mount_gradio_app(app, demo, path="/", footer_links=[], max_file_size="25mb", ssr_mode=False,
+                             css=(ROOT / "studio.css").read_text(), js=(ROOT / "studio.js").read_text(),
+                             theme=gr.themes.Base(primary_hue="neutral", secondary_hue="neutral", neutral_hue="stone", radius_size="lg"))
+    return app

hero.html

@@ -0,0 +1,13 @@
+<header class="hero">
+    <h1>🎭 Drama 3</h1>
+    <p>by Fish Audio · served with SGLang Omni</p>
+    <p>Write a line, pick a voice, and <strong>direct the performance</strong>: drop in a <code>[laughs]</code>, whisper a phrase, or describe a voice that doesn't exist yet. Then give the actor a note and get another take. The 35B MoE streams from a single GPU as it generates.</p>
+    <nav class="links resource-links" aria-label="Drama 3 resources">
+      <a href="https://sgl-project.github.io/sglang-omni/cookbook/fishaudio_drama3.html" target="_blank" rel="noreferrer">SGLang Omni cookbook</a>
+      <a href="https://github.com/sgl-project/sglang-omni" target="_blank" rel="noreferrer">⭐ Star SGLang Omni</a>
+      <a href="https://github.com/fishaudio/fish-speech/tree/main/docs/assets/Fish_Audio_Drama_3_Technical_Report.pdf" target="_blank" rel="noreferrer">Technical report</a>
+      <a href="https://fish.audio" target="_blank" rel="noreferrer">Fish Audio</a>
+      <a href="https://github.com/fishaudio/fish-speech" target="_blank" rel="noreferrer">Fish Speech</a>
+      <a href="https://huggingface.co/fishaudio/drama-3" target="_blank" rel="noreferrer">Model card</a>
+    </nav>
+</header>

index.html

@@ -0,0 +1,1210 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="UTF-8">
+<meta name="viewport" content="width=device-width, initial-scale=1.0">
+<title>Fish Audio Drama 3</title>
+<style>
+*, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; }
+
+:root {
+  --bg: #f6f5f3;
+  --card: #ffffff;
+  --soft: #efede9;
+  --field: #ffffff;
+  --ink: #141414;
+  --muted: #6b6b66;
+  --faint: #a3a19b;
+  --border: #e2e0dc;
+  --card-border: #e8e6e2;
+  --accent: #141414;
+  --accent-ink: #f6f5f3;
+  --accent-soft: rgba(20, 20, 20, 0.08);
+  --moment-bg: #fbe9c4; --moment-ink: #7d4b00;
+  --scope-bg: #d9ebfa; --scope-ink: #0b5687;
+  --mood-bg: #ebe1fa; --mood-ink: #56289e;
+  --note-bg: #e3f1e3; --note-ink: #2a662d;
+  --ok: #1f9d55;
+  --warn: #b7791f;
+  --err: #d64545;
+  --hover: brightness(0.96);
+  --shadow: 0 1px 2px rgba(20, 20, 20, 0.04);
+  --radius: 16px;
+  --sans: ui-sans-serif, -apple-system, "Segoe UI", Inter, Roboto, sans-serif;
+  --mono: ui-monospace, "SF Mono", Menlo, Consolas, monospace;
+  color-scheme: light;
+}
+@media (prefers-color-scheme: dark) {
+  :root:not([data-theme="light"]) {
+    --bg: #161615; --card: #201f1d; --soft: #2a2826; --field: #1a1918; --ink: #f3f2ef; --muted: #a8a7a1;
+    --faint: #6f6d68; --border: #35332f; --card-border: #35332f;
+    --accent: #f3f2ef; --accent-ink: #161615; --accent-soft: rgba(243, 242, 239, 0.12);
+    --moment-bg: #463312; --moment-ink: #ffd68a;
+    --scope-bg: #133852; --scope-ink: #a3d4fa;
+    --mood-bg: #37245b; --mood-ink: #d9c4ff;
+    --note-bg: #1e3a20; --note-ink: #b5e3b7;
+    --hover: brightness(1.15);
+    --shadow: 0 1px 2px rgba(0, 0, 0, 0.3); color-scheme: dark;
+  }
+}
+:root[data-theme="dark"] {
+  --bg: #161615; --card: #201f1d; --soft: #2a2826; --field: #1a1918; --ink: #f3f2ef; --muted: #a8a7a1;
+  --faint: #6f6d68; --border: #35332f; --card-border: #35332f;
+  --accent: #f3f2ef; --accent-ink: #161615; --accent-soft: rgba(243, 242, 239, 0.12);
+  --moment-bg: #463312; --moment-ink: #ffd68a;
+  --scope-bg: #133852; --scope-ink: #a3d4fa;
+  --mood-bg: #37245b; --mood-ink: #d9c4ff;
+  --note-bg: #1e3a20; --note-ink: #b5e3b7;
+  --hover: brightness(1.15);
+  --shadow: 0 1px 2px rgba(0, 0, 0, 0.3); color-scheme: dark;
+}
+
+html { -webkit-text-size-adjust: 100%; }
+body {
+  background: var(--bg); color: var(--ink);
+  font-family: var(--sans);
+  font-size: 14px; line-height: 1.5;
+}
+button, input, textarea, select { font: inherit; color: inherit; }
+button { cursor: pointer; }
+a { color: inherit; }
+code { font-family: var(--mono); font-size: .92em; background: var(--soft); padding: 0 4px; border-radius: 5px; }
+
+.app { max-width: 980px; margin: 0 auto; padding: 26px 32px 48px; display: flex; flex-direction: column; gap: 16px; }
+
+.hero { margin: 36px 0 8px; }
+.hero-top { display: flex; align-items: center; gap: 8px; flex-wrap: wrap; }
+.hero h1 {
+  font-size: clamp(30px, 5vw, 46px); font-weight: 800;
+  letter-spacing: -0.025em; line-height: 1.08;
+  max-width: 640px; margin: 18px 0 10px; color: var(--ink);
+}
+.hero p { font-size: 15px; line-height: 1.55; color: var(--muted); max-width: 640px; }
+.hero p + p { margin-top: 10px; }
+.hero p strong { color: var(--ink); font-weight: 600; }
+.links { display: flex; gap: 8px; flex-wrap: wrap; }
+.links a {
+  border: 1px solid var(--border); border-radius: 999px;
+  background: var(--card); padding: 6px 14px; line-height: 1.4;
+  font-size: 13px; font-weight: 500; color: var(--ink); text-decoration: none;
+}
+.links a:hover { border-color: var(--ink); }
+.spacer { flex: 1; }
+.pill {
+  display: inline-flex; align-items: center; gap: 7px; max-width: 100%;
+  border: 0; border-radius: 999px; background: var(--card); color: var(--ink);
+  font-size: 13px; font-weight: 500; line-height: 1.4; padding: 5px 14px; white-space: nowrap; overflow: hidden;
+}
+.pill b { color: var(--ink); font-weight: 600; }
+.dot { width: 7px; height: 7px; border-radius: 50%; background: var(--faint); flex-shrink: 0; }
+.dot.ready { background: var(--ok); box-shadow: 0 0 0 3px rgba(31, 157, 85, 0.15); }
+.dot.busy { background: var(--warn); animation: blink 1s infinite; }
+.dot.err { background: var(--err); }
+.account { position: relative; }
+.pill.signin { text-decoration: none; font-weight: 600; }
+.pill.signin:hover { box-shadow: inset 0 0 0 1px var(--border); }
+.msg a { color: inherit; font-weight: 700; margin-left: 6px; }
+.account .pill { cursor: pointer; font-family: inherit; }
+.account .pill:hover { box-shadow: inset 0 0 0 1px var(--border); }
+.menu {
+  position: absolute; right: 0; top: calc(100% + 6px); z-index: 40; min-width: 210px;
+  background: var(--card); border: 1px solid var(--card-border); border-radius: 12px; padding: 6px;
+  box-shadow: 0 12px 32px rgba(0, 0, 0, 0.18);
+}
+.menu[hidden] { display: none; }
+.menu .sum { padding: 6px 8px 8px; font-size: 12px; color: var(--muted); border-bottom: 1px solid var(--border); margin-bottom: 4px; }
+.menu a { display: block; padding: 7px 8px; border-radius: 8px; font-size: 13px; text-decoration: none; }
+.menu a:hover { background: var(--soft); }
+.ibtn {
+  width: 28px; height: 28px; display: inline-flex; align-items: center; justify-content: center;
+  border: 0; border-radius: 999px; background: var(--card); color: var(--ink);
+}
+.ibtn:hover { box-shadow: inset 0 0 0 1px var(--border); }
+.ibtn svg { width: 14px; height: 14px; }
+
+.card { background: var(--card); border: 1px solid var(--card-border); border-radius: var(--radius); box-shadow: var(--shadow); padding: 14px; }
+.card-head { display: flex; align-items: center; gap: 8px; margin-bottom: 10px; padding: 0 2px; }
+.card-head h2 { font-size: 14px; font-weight: 600; color: var(--ink); }
+.card-head .spacer { flex: 1; }
+.linkbtn {
+  border: 1px solid var(--border); border-radius: 999px; background: transparent; color: var(--muted);
+  padding: 3px 11px; font-size: 12.5px; font-weight: 500; transition: color .12s ease, border-color .12s ease;
+}
+.linkbtn:hover { color: var(--ink); border-color: var(--ink); }
+
+.boot { display: none; border-color: var(--border); background: var(--card); }
+.boot.show { display: block; }
+.boot h3 { font-size: 15px; font-weight: 700; letter-spacing: -0.01em; margin-bottom: 2px; }
+.boot p { color: var(--muted); font-size: 13px; }
+.boot .log { margin-top: 8px; font-family: var(--mono); font-size: 11px; color: var(--muted); white-space: nowrap; overflow: hidden; text-overflow: ellipsis; }
+.boot .bar { height: 3px; border-radius: 2px; background: var(--soft); overflow: hidden; margin-top: 10px; }
+.boot .bar i { display: block; height: 100%; width: 30%; background: var(--ink); border-radius: 2px; animation: sweep 1.6s ease-in-out infinite; }
+.boot.error { border-color: var(--err); }
+.boot.error .bar { display: none; }
+
+.editor {
+  position: relative; border: 1px solid var(--border); border-radius: 12px;
+  background: var(--field); transition: border-color .12s ease;
+}
+.editor:focus-within { border-color: var(--muted); }
+.mirror, textarea.script {
+  font-family: var(--sans); font-size: 16px; line-height: 1.6; letter-spacing: normal; font-weight: 400;
+  padding: 12px 14px; margin: 0; border: 0;
+  white-space: pre-wrap; overflow-wrap: break-word; word-break: normal;
+  tab-size: 4; box-sizing: border-box;
+}
+.mirror { min-height: 120px; color: var(--ink); pointer-events: none; }
+textarea.script {
+  position: absolute; inset: 0; width: 100%; height: 100%;
+  resize: none; overflow: hidden; background: transparent;
+  color: transparent; caret-color: var(--ink); outline: none;
+}
+textarea.script::placeholder { color: var(--muted); opacity: .75; }
+textarea.script::selection { background: rgba(120, 120, 255, .28); color: transparent; }
+.mirror mark { color: inherit; border-radius: 4px; padding: 0; margin: 0; }
+.mirror .h-point { background: var(--moment-bg); color: var(--moment-ink); }
+.mirror .h-tag { background: var(--scope-bg); color: var(--scope-ink); }
+.mirror .h-span { text-decoration: underline; text-decoration-color: var(--scope-ink); text-decoration-thickness: 2px; text-underline-offset: 4px; }
+.mirror .h-style { background: var(--mood-bg); color: var(--mood-ink); }
+.mirror .h-note { background: var(--note-bg); color: var(--note-ink); }
+.field:focus { border-color: var(--muted); }
+.script-meta { display: flex; justify-content: space-between; align-items: center; margin-top: 6px; padding: 0 2px; font-size: 12px; color: var(--muted); gap: 8px; }
+#charCount { font-variant-numeric: tabular-nums; white-space: nowrap; }
+.palette { display: grid; grid-template-columns: minmax(0, 1fr); gap: 7px; margin-top: 12px; padding: 0 2px; }
+.prow { display: flex; align-items: center; gap: 5px; flex-wrap: wrap; }
+.plabel {
+  display: inline-flex; align-items: center; gap: 5px; height: 25px; width: 72px; flex-shrink: 0;
+  font-size: 11px; font-weight: 700; letter-spacing: .04em; text-transform: uppercase; color: var(--muted);
+}
+.plabel i { width: 7px; height: 7px; border-radius: 50%; display: inline-block; flex: none; }
+.prow > .prow { flex: 1; min-width: 0; }
+.chip {
+  display: inline-flex; align-items: center; gap: 4px; height: 27px; box-sizing: border-box;
+  border: 1px solid var(--border); background: var(--card); border-radius: 999px;
+  padding: 0 11px; font-size: 12.5px; font-weight: 500; line-height: 1; white-space: nowrap;
+  transition: border-color .12s ease, background .12s ease, filter .12s ease;
+}
+.chip:hover { border-color: var(--ink); }
+.chip.on { background: var(--ink); color: var(--bg); border-color: var(--ink); }
+.chip.dashed { border-style: dashed; background: transparent; color: var(--muted); }
+.chip.dashed:hover { color: var(--ink); border-color: var(--muted); }
+.chip:disabled { opacity: .5; cursor: default; }
+.palette .chip { height: 25px; padding: 0 9px; border-color: transparent; }
+.palette .chip:hover { filter: var(--hover); border-color: transparent; }
+.prow[data-kind=point] .plabel i { background: var(--moment-ink); }
+.prow[data-kind=scope] .plabel i { background: var(--scope-ink); }
+.prow[data-kind=style] .plabel i { background: var(--mood-ink); }
+.prow[data-kind=point] .chip { background: var(--moment-bg); color: var(--moment-ink); }
+.prow[data-kind=scope] .chip { background: var(--scope-bg); color: var(--scope-ink); }
+.prow[data-kind=style] .chip { background: var(--mood-bg); color: var(--mood-ink); }
+.examples { display: flex; gap: 6px; overflow-x: auto; padding: 0 2px 2px; scrollbar-width: none; }
+.examples::-webkit-scrollbar { display: none; }
+
+.seg { display: flex; background: var(--soft); border-radius: 999px; padding: 3px; gap: 3px; margin-bottom: 12px; }
+.seg button { flex: 1; border: none; background: none; border-radius: 999px; padding: 7px 6px; font-weight: 600; font-size: 13px; color: var(--muted); }
+.seg button.on { background: var(--card); color: var(--ink); box-shadow: 0 1px 2px rgba(20, 20, 20, 0.08); }
+.pane { display: none; flex-direction: column; gap: 10px; }
+.pane.on { display: flex; }
+.hint { font-size: 12.5px; color: var(--muted); }
+.hint b { color: var(--ink); font-weight: 600; }
+.voices { display: flex; gap: 6px; flex-wrap: wrap; align-items: center; }
+.rec { display: inline-flex; align-items: center; gap: 6px; }
+.rec.live .chip { border-color: var(--err); color: var(--err); }
+.meter { width: 70px; height: 4px; background: var(--soft); border-radius: 2px; overflow: hidden; display: none; }
+.rec.live .meter { display: block; }
+.meter i { display: block; height: 100%; width: 0; background: var(--err); }
+.preview { width: 100%; height: 34px; }
+label.fl { display: block; font-size: 13px; font-weight: 600; margin-bottom: 5px; }
+label.fl span { font-weight: 400; color: var(--muted); }
+.field {
+  width: 100%; border: 1px solid var(--border); border-radius: 12px; background: var(--field);
+  padding: 9px 12px; font-size: 14px; outline: none; resize: vertical; transition: border-color .12s ease;
+}
+.row2 { display: grid; grid-template-columns: 1fr 2fr; gap: 10px; }
+.badge { display: inline-block; font-size: 10.5px; font-weight: 600; letter-spacing: .03em; padding: 1px 7px; border-radius: 999px; background: var(--scope-bg); color: var(--scope-ink); vertical-align: 1px; }
+
+.perform-row { display: flex; gap: 8px; align-items: center; }
+.perform {
+  flex: 1; display: inline-flex; align-items: center; justify-content: center; gap: 10px; min-height: 46px;
+  border: none; border-radius: 999px; background: var(--ink); color: var(--bg);
+  padding: 12px 20px; font-size: 15px; font-weight: 600; transition: transform .08s, box-shadow .2s, opacity .12s;
+}
+.perform:hover:not(:disabled) { box-shadow: 0 0 0 4px var(--accent-soft); }
+.perform:active:not(:disabled) { transform: scale(0.99); }
+.perform:disabled { opacity: .45; cursor: default; }
+.perform.live:disabled { opacity: 1; }
+.perform.live { background: var(--ink); color: var(--bg); box-shadow: 0 0 0 4px var(--accent-soft); }
+.perform kbd { font-family: var(--mono); font-size: 11px; opacity: .55; font-weight: 500; }
+.onair { width: 8px; height: 8px; border-radius: 50%; background: var(--err); display: none; animation: blink .9s infinite; }
+.perform.live .onair { display: inline-block; }
+.stop { border: 1px solid var(--border); background: var(--card); border-radius: 999px; min-height: 46px; padding: 0 18px; font-weight: 600; display: none; }
+.stop:hover { border-color: var(--ink); }
+.stop.show { display: inline-block; }
+
+.msg { display: none; padding: 9px 12px; border-radius: 12px; font-size: 13px; }
+.msg.show { display: block; }
+.msg.err { background: rgba(214, 69, 69, 0.08); color: var(--err); border: 1px solid rgba(214, 69, 69, 0.25); }
+.msg.info { background: var(--card); color: var(--muted); border: 1px solid var(--card-border); }
+
+.stage { display: none; }
+.stage.show { display: block; }
+canvas.viz { width: 100%; height: 76px; display: block; border-radius: 12px; background: var(--field); border: 1px solid var(--border); }
+.metrics { display: flex; flex-wrap: wrap; gap: 6px 14px; margin-top: 10px; font-size: 12px; color: var(--muted); }
+.metrics b { font-family: var(--mono); color: var(--ink); font-weight: 600; }
+.metrics .hot b { color: var(--ink); }
+.take-tools { display: flex; gap: 6px; flex-wrap: wrap; align-items: center; margin-top: 10px; }
+.take-tools audio { flex: 1 1 220px; height: 34px; min-width: 0; }
+.direct { margin-top: 14px; padding-top: 14px; border-top: 1px solid var(--card-border); }
+.direct h3 { font-size: 14px; font-weight: 600; margin-bottom: 2px; }
+.direct .hint { margin-bottom: 8px; }
+.note-row { display: flex; gap: 6px; margin-top: 8px; }
+.note-row .field { flex: 1; min-width: 0; border-radius: 999px; padding: 8px 14px; }
+.btn {
+  border: 1px solid var(--ink); background: var(--ink); color: var(--bg); border-radius: 999px;
+  padding: 8px 16px; font-weight: 600; font-size: 13px; white-space: nowrap;
+}
+.btn.ghost { background: var(--card); color: var(--ink); border-color: var(--border); }
+.btn.ghost:hover { border-color: var(--ink); }
+.btn:disabled { opacity: .45; cursor: default; }
+
+.slate { display: flex; flex-direction: column; gap: 6px; margin-top: 12px; }
+.slate-item {
+  display: grid; grid-template-columns: auto 1fr auto; align-items: center; gap: 10px;
+  padding: 8px 10px; border: 1px solid var(--border); border-radius: 12px; background: var(--field);
+}
+.slate-item.current { border-color: var(--ink); }
+.slate-num { font-family: var(--mono); font-size: 11px; color: var(--faint); }
+.slate-body { min-width: 0; }
+.slate-title { font-size: 13px; font-weight: 600; white-space: nowrap; overflow: hidden; text-overflow: ellipsis; }
+.slate-sub { font-size: 11.5px; color: var(--muted); white-space: nowrap; overflow: hidden; text-overflow: ellipsis; }
+.slate-actions { display: flex; gap: 4px; }
+.mini { border: 1px solid var(--border); background: var(--card); border-radius: 999px; padding: 3px 10px; font-size: 12px; font-weight: 500; }
+.mini:hover { border-color: var(--ink); }
+
+.sheet-bg { position: fixed; inset: 0; background: rgba(0,0,0,.45); display: none; align-items: center; justify-content: center; padding: 16px; z-index: 50; }
+.sheet-bg.open { display: flex; }
+.sheet { background: var(--card); border: 1px solid var(--card-border); border-radius: 18px; width: 100%; max-width: 560px; max-height: 86vh; overflow: auto; padding: 18px; }
+.sheet h3 { font-size: 16px; font-weight: 700; letter-spacing: -0.01em; margin-bottom: 4px; }
+.sheet > p { color: var(--muted); font-size: 13px; margin-bottom: 12px; }
+.cheat { width: 100%; border-collapse: collapse; font-size: 13px; }
+.cheat td { padding: 8px 6px; border-top: 1px solid var(--card-border); vertical-align: top; }
+.cheat td:first-child { font-weight: 600; white-space: nowrap; }
+.cheat code { font-size: 12px; padding: 1px 5px; }
+.srow { display: grid; grid-template-columns: 110px 1fr 44px; gap: 10px; align-items: center; margin: 10px 0; font-size: 13px; }
+.srow input[type=range] { width: 100%; accent-color: var(--ink); }
+.srow output { font-family: var(--mono); font-size: 12px; text-align: right; color: var(--muted); }
+
+details.credits { color: var(--muted); font-size: 13px; }
+details.credits summary { cursor: pointer; color: var(--ink); font-weight: 600; font-size: 14px; }
+details.credits p { margin-top: 8px; }
+.fine { color: var(--muted); font-size: 12px; }
+
+@media (max-width: 560px) {
+  .row2 { grid-template-columns: 1fr; }
+  .plabel { width: 100%; }
+  .palette .prow > .prow { flex-wrap: nowrap; overflow-x: auto; scrollbar-width: none; padding-bottom: 2px; }
+  .palette .prow > .prow::-webkit-scrollbar { display: none; }
+  .app { padding: 10px 16px 48px; }
+  .hero { margin-top: 20px; }
+  .perform kbd { display: none; }
+  .srow { grid-template-columns: 90px 1fr 40px; }
+}
+
+@keyframes blink { 0%, 100% { opacity: 1; } 50% { opacity: .25; } }
+@keyframes sweep { 0% { transform: translateX(-110%); } 100% { transform: translateX(360%); } }
+
+/* Keep the resources in the introduction, before the studio. */
+.resource-links { gap: 8px; margin: 0 0 6px; }
+.resource-links a { display: inline-flex; align-items: center; min-height: 40px; padding: 9px 16px; background: var(--card); font-size: 14px; font-weight: 600; }
+.credits { font-size: 13px; }
+.credits h2 { margin: 0 0 10px; color: var(--ink); font-size: 16px; font-weight: 650; }
+@media (max-width: 640px) {
+  .resource-links { gap: 7px; }
+  .resource-links a { min-height: 40px; padding: 8px 12px; font-size: 13px; }
+}
+
+.direction-dropdown { margin-top: 12px; padding: 12px; border: 1px solid var(--border); border-radius: 12px; }
+.direction-dropdown > summary { cursor: pointer; font-size: 13px; font-weight: 600; }
+
+</style>
+</head>
+<body>
+<div class="app">
+
+  <header class="hero">
+    <div class="hero-top">
+      <span class="pill" id="statusPill"><span class="dot" id="statusDot"></span><span id="statusText">connecting…</span></span>
+      <span class="spacer"></span>
+      <a class="pill signin" id="signIn" href="/oauth/huggingface/login" target="_blank" rel="noopener" hidden></a>
+      <div class="account" id="account" hidden>
+        <button class="pill" id="userPill" type="button" aria-haspopup="true" aria-expanded="false" onclick="toggleAccount(event)"></button>
+        <div class="menu" id="accountMenu" hidden>
+          <div class="sum" id="accountSummary"></div>
+          <a href="https://huggingface.co/pricing" target="_blank" rel="noreferrer">Unlimited takes with PRO</a>
+          <a href="/oauth/huggingface/logout">Sign out</a>
+        </div>
+      </div>
+      <button class="ibtn" type="button" onclick="openSheet('settingsSheet')" title="Sampling settings" aria-label="Sampling settings">
+        <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round"><line x1="4" y1="21" x2="4" y2="14"/><line x1="4" y1="10" x2="4" y2="3"/><line x1="12" y1="21" x2="12" y2="12"/><line x1="12" y1="8" x2="12" y2="3"/><line x1="20" y1="21" x2="20" y2="16"/><line x1="20" y1="12" x2="20" y2="3"/><line x1="1" y1="14" x2="7" y2="14"/><line x1="9" y1="8" x2="15" y2="8"/><line x1="17" y1="16" x2="23" y2="16"/></svg>
+      </button>
+      <button class="ibtn" type="button" onclick="toggleTheme()" title="Toggle theme" aria-label="Toggle theme">
+        <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round"><path d="M21 12.79A9 9 0 1 1 11.21 3 7 7 0 0 0 21 12.79z"/></svg>
+      </button>
+    </div>
+    <h1>🎭 Drama 3</h1>
+    <p>by Fish Audio · served with SGLang Omni</p>
+    <p>Write a line, pick a voice, and <strong>direct the performance</strong>: drop in a <code>[laughs]</code>, whisper a phrase, or describe a voice that doesn't exist yet. Then give the actor a note and get another take. The 35B MoE streams from a single GPU as it generates.</p>
+    <nav class="links resource-links" aria-label="Drama 3 resources">
+      <a href="https://sgl-project.github.io/sglang-omni/cookbook/fishaudio_drama3.html" target="_blank" rel="noreferrer">SGLang Omni cookbook</a>
+      <a href="https://github.com/sgl-project/sglang-omni" target="_blank" rel="noreferrer">⭐ Star SGLang Omni</a>
+      <a href="https://github.com/fishaudio/fish-speech/tree/main/docs/assets/Fish_Audio_Drama_3_Technical_Report.pdf" target="_blank" rel="noreferrer">Technical report</a>
+      <a href="https://fish.audio" target="_blank" rel="noreferrer">Fish Audio</a>
+      <a href="https://github.com/fishaudio/fish-speech" target="_blank" rel="noreferrer">Fish Speech</a>
+      <a href="https://huggingface.co/fishaudio/drama-3" target="_blank" rel="noreferrer">Model card</a>
+    </nav>
+  </header>
+
+  <section class="card boot" id="boot">
+    <h3 id="bootTitle">Setting the stage…</h3>
+    <p id="bootPhase">Launching SGLang Omni</p>
+    <div class="bar"><i></i></div>
+    <div class="log" id="bootLog"></div>
+  </section>
+
+  <!-- Script -->
+  <section class="card">
+    <div class="card-head">
+      <h2>Script</h2>
+      <span class="spacer"></span>
+      <button class="linkbtn" type="button" onclick="openSheet('helpSheet')">How to direct</button>
+    </div>
+    <div class="examples" id="examples"></div>
+    <div class="editor" style="margin-top:10px"><div class="mirror" id="mirror" aria-hidden="true"></div><textarea class="script" id="script" spellcheck="true" aria-label="Script"></textarea></div>
+    <details class="direction-dropdown">
+      <summary>Direction tags</summary>
+      <div class="script-meta"><span>Tags in <code>[ ]</code> and <code>&lt; &gt;</code> are direction, not words. Any wording works.</span><span id="charCount">0</span></div>
+      <div class="palette">
+        <div class="prow" data-kind="point"><span class="plabel"><i></i>Events</span><span id="pointTags" class="prow"></span></div>
+        <div class="prow" data-kind="scope"><span class="plabel"><i></i>Span</span><span id="scopedTags" class="prow"></span></div>
+        <div class="prow" data-kind="style"><span class="plabel"><i></i>Mood</span><span id="styleTags" class="prow"></span></div>
+      </div>
+    </details>
+  </section>
+
+  <!-- Voice -->
+  <section class="card">
+    <div class="card-head"><h2>Voice</h2></div>
+    <div class="seg" role="tablist">
+      <button type="button" data-mode="random" onclick="setMode('random')">🎲 Surprise</button>
+      <button type="button" data-mode="clone" onclick="setMode('clone')">🎙 Clone</button>
+      <button type="button" data-mode="design" onclick="setMode('design')">✨ Design</button>
+    </div>
+
+    <div class="pane" id="pane-clone">
+      <div class="voices" id="voiceChips"></div>
+      <audio class="preview" id="voicePreview" controls preload="none"></audio>
+      <div>
+        <label class="fl" for="transcript">What the clip says <span id="transcriptState"></span></label>
+        <textarea class="field" id="transcript" rows="2" placeholder="Exact words spoken in the clip"></textarea>
+      </div>
+      <p class="hint">Upload or record a few seconds of one speaker. Clips are trimmed to 20 s and transcribed for you. Only clone voices you have permission to use.</p>
+    </div>
+
+    <div class="pane" id="pane-design">
+      <div class="voices" id="castChips"></div>
+      <div class="row2">
+        <div><label class="fl" for="charName">Character</label><input class="field" id="charName" maxlength="40" placeholder="Captain"></div>
+        <div><label class="fl" for="charDesc">Voice description</label><textarea class="field" id="charDesc" rows="2" maxlength="300" placeholder="gravelly old male voice, about seventy, weathered and slow"></textarea></div>
+      </div>
+      <div><label class="fl" for="charScene">Scene <span>(optional)</span></label><input class="field" id="charScene" maxlength="200" placeholder="a storm-battered tavern at the edge of the harbor"></div>
+      <p class="hint">No audio needed: Drama 3 casts the voice from your description, the way it voices characters in a scene. Describe age, timbre, accent, and energy. Like a take? <b>Keep this voice</b> turns it into a clone.</p>
+    </div>
+
+    <div class="pane" id="pane-random">
+      <p class="hint">No reference and no description: Drama 3 improvises a voice that fits the line. Every take is a stranger, so hit <b>Keep this voice</b> when you meet one you like.</p>
+    </div>
+  </section>
+
+  <div class="perform-row">
+    <button class="perform" id="performBtn" type="button" onclick="perform()" disabled>
+      <span class="onair"></span><span id="performLabel">Warming up…</span><kbd>⌘↵</kbd>
+    </button>
+    <button class="stop" id="stopBtn" type="button" onclick="stopTake()">Stop</button>
+  </div>
+  <div class="msg" id="msg"></div>
+
+  <!-- Stage -->
+  <section class="card stage" id="stage">
+    <div class="card-head"><h2 id="takeTitle">Take 1</h2><span class="spacer"></span><span class="hint" id="takeVoice"></span></div>
+    <canvas class="viz" id="viz"></canvas>
+    <div class="metrics" id="metrics"></div>
+    <div class="take-tools">
+      <audio id="takeAudio" controls preload="auto"></audio>
+      <button class="btn ghost" type="button" id="downloadBtn" onclick="downloadTake()" disabled>Download</button>
+      <button class="btn ghost" type="button" id="keepBtn" onclick="keepVoice()" disabled title="Use this take as a clone reference">Keep this voice</button>
+    </div>
+    <div class="direct" id="direct">
+      <h3>🎬 Give a note</h3>
+      <p class="hint">The actor hears this take, then redoes the line with your note, like the multi-turn directing in the Drama 3 paper.</p>
+      <div class="voices" id="noteChips"></div>
+      <div class="note-row">
+        <input class="field" id="noteIn" maxlength="120" placeholder="e.g. more amused than before; drop the hurt">
+        <button class="btn" type="button" id="retakeBtn" onclick="retake()">Retake</button>
+      </div>
+    </div>
+    <div class="slate" id="slate"></div>
+  </section>
+
+  <!-- Acknowledgments are always visible. -->
+  <section class="card">
+
+    <section class="credits" aria-labelledby="acknowledgments-title">
+      <h2 id="acknowledgments-title">Acknowledgments</h2>
+      <p>It took us over half a year to bring SGLang Omni back to life. Thanks to all the people who helped us along the way.</p>
+      <p>Thanks so much to <a href="https://huggingface.co/andito" target="_blank" rel="noreferrer">andito</a>, <a href="https://x.com/_akhaliq" target="_blank" rel="noreferrer">AK</a>, and <a href="https://huggingface.co/multimodalart" target="_blank" rel="noreferrer">poli</a> for helping set up the Space.</p>
+      <p><b>SGLang Omni team:</b> <a href="https://x.com/GenAI_is_real" target="_blank" rel="noreferrer">Chenyang Zhao</a>, Jingwen Gu, Shenggui Li, Richard Wang, Yichi Zhang, Xuesong Ye, Jiaxin Deng, Ratish P, Jiaxuan Luo, Xinyu Lu, Wenyao Gao, Ao Sun, Jintao Quaige Li, Yuhao Chen, Yuhan Zhou, Xuxiang Tian, Ruoyi Zhang, Junnan Li, Yixiao Zeng, Chenchen Hong, Huapeng Zhou, Yang Gao, Chaoye Fei, Yiyang Zh, Kang, Yitian Gong, Yuan Luo, Hao Jin, Yifei Gao, Gaokai Zhang, Lifan Shen, Hui Shen, Kejing Xia, Jinxin Li, Dongming Shen, Jun Liu, Yue Yin, Zhihao Guo, Jisen Li, Ruoyu He, Xuanyi Li, Chenye Zhu, Liang Zhen, Xiaodong Ye, Yizhi Zhen, Kaizuo Li, Chonghan Liu, Wei Guan, Zhuofeng Li</p>
+      <p><b>FishAudio Team:</b> Shijia Liao, Shidong Li, Xuehao Yang, Jiahua Liu</p>
+    </section>
+    <p class="fine" style="margin-top:10px">Model license: <a href="https://huggingface.co/fishaudio/drama-3/blob/main/LICENSE" target="_blank" rel="noreferrer">Fish Audio Community License</a>, for research and non-commercial use. Takes and uploaded voices live in server memory only and expire.</p>
+  </section>
+</div>
+
+<!-- Help -->
+<div class="sheet-bg" id="helpSheet" onclick="closeSheet(event)">
+  <div class="sheet" role="dialog" aria-label="How to direct">
+    <h3>How to direct Drama 3</h3>
+    <p>Direction is plain language inside the script. Brackets and angle tags are never spoken; round parentheses are. These are the controls from the technical report (Appendix C); the wording is yours.</p>
+    <table class="cheat">
+      <tr><td>Point tag</td><td>An event at that spot.<br><code>Well [sighs] fine, you win. [laughs]</code></td></tr>
+      <tr><td>Scoped tag</td><td>Changes only the words inside.<br><code>Promise me &lt;whisper&gt;you won't tell her&lt;/whisper&gt;.</code></td></tr>
+      <tr><td>Style prefix</td><td>Sets the delivery of the whole line.<br><code>[icy calm, as if you already accepted it] Fine.</code></td></tr>
+      <tr><td>Director's note</td><td>Intent, pacing, emphasis, on two lines.<br><code>Direction: put the rise on "goodbye".<br>Line: So this is goodbye.</code></td></tr>
+      <tr><td>Retake note</td><td>After a take, a note relative to it.<br><code>[softer and warmer versus previous]</code></td></tr>
+      <tr><td>Character</td><td>In ✨ Design, a description casts the voice.<br><code>Emma: young female voice, sleepy, half talking to herself.</code></td></tr>
+    </table>
+    <div style="margin-top:14px;text-align:right"><button class="btn" type="button" onclick="closeSheet()">Got it</button></div>
+  </div>
+</div>
+
+<!-- Settings -->
+<div class="sheet-bg" id="settingsSheet" onclick="closeSheet(event)">
+  <div class="sheet" role="dialog" aria-label="Sampling settings">
+    <h3>Sampling</h3>
+    <p>Defaults: temperature 1.0, top-p 0.9, and no top-k cap. A low top-k makes takes go silent mid-line, so 0 (no cap) is recommended.</p>
+    <div class="srow"><label for="temp">Temperature</label><input type="range" id="temp" min="0.3" max="1.3" step="0.05"><output id="tempV"></output></div>
+    <div class="srow"><label for="topp">Top-p</label><input type="range" id="topp" min="0.5" max="1" step="0.01"><output id="toppV"></output></div>
+    <div class="srow"><label for="topk">Top-k</label><input type="range" id="topk" min="0" max="100" step="1"><output id="topkV"></output></div>
+    <div style="display:flex;gap:8px;justify-content:flex-end;margin-top:14px">
+      <button class="btn ghost" type="button" onclick="resetSampling()">Reset</button>
+      <button class="btn" type="button" onclick="closeSheet()">Done</button>
+    </div>
+  </div>
+</div>
+
+<input type="file" id="fileIn" accept="audio/*,.wav,.mp3,.m4a,.flac,.ogg,.webm" hidden>
+
+<script>
+// ── Content ─────────────────────────────────────────────────────────────────
+const EXAMPLES = [
+  {
+    "label": "Stage fright",
+    "text": "[clears throat] Ladies and gentlemen... [hesitates] I had a speech. [nervous laugh] I really did have a speech."
+  },
+  {
+    "label": "We won!",
+    "text": "[with bright eager excitement] We won! We actually won! [laughs] I can't believe it!"
+  },
+  {
+    "label": "Missed train",
+    "text": "[sighs] Fine. If we miss it, it's on you. [laughs] Well, maybe that's fair."
+  },
+  {
+    "label": "A secret",
+    "text": "Promise me <whisper>you won't tell her</whisper>. [pause] Not a word. Not even a hint."
+  },
+  {
+    "label": "Bad news",
+    "text": "[with heavy, deflated sadness] The results just came in. We didn't make it.",
+    "note": "more amused than before; drop the hurt"
+  },
+  {
+    "label": "Goodbye",
+    "text": "Direction: icy calm, as if you already accepted it. Put the rise on \"goodbye\".\nLine: So this is goodbye."
+  },
+  {
+    "label": "Commentator",
+    "text": "[breathless, rising excitement] He's past one... past two... he shoots, and [shouting] HE SCORES! Unbelievable!"
+  },
+  {
+    "label": "Museum guard",
+    "text": "Direction: deadpan, a bored museum guard at the end of a very long shift.\nLine: Please do not touch the dinosaur. Again."
+  }
+];
+const CAST = [
+  { label: '🏴‍☠️ Sea captain', name: 'Captain', desc: 'gravelly old male voice, about seventy, weathered and slow, salt in every word',
+    scene: 'a storm-battered tavern at the edge of the harbor', text: "[low and menacing] Ye want to know what's out past the reef? [pause] Pray ye never find out." },
+  { label: '📻 Night radio', name: 'Night Owl', desc: 'smooth warm baritone, a late-night radio host, relaxed and intimate, close to the mic',
+    scene: '', text: "It's two in the morning, the city's finally asleep... [chuckles] and you're still here with me. Good." },
+  { label: '🧒 Hyper kid', name: 'Mia', desc: 'excited little girl, about seven, talks fast, giggly and breathless',
+    scene: '', text: '[gasps] Did you see that? The dog just ate a WHOLE pizza! [laughs] Can we get a dog? Please please please?' },
+  { label: '🕵️ Noir detective', name: 'Detective', desc: 'world-weary man in his forties, low smoky voice, clipped and dry',
+    scene: 'a rain-soaked office, one desk lamp still on', text: 'She walked in at midnight. [pause] Trouble usually knocks. This one had a key.' },
+  { label: '📯 Royal herald', name: 'Herald', desc: 'booming, pompous royal herald, theatrical and grandly formal',
+    scene: '', text: '[clears throat] Hear ye, hear ye! By order of the Queen, all cheese is hereby... [hesitates] mandatory.' },
+  { label: '😴 Sleepy intern', name: 'Emma', desc: 'young female voice, sleepy, half talking to herself',
+    scene: 'a small office on a Monday morning; the weekly meeting is about to start', text: '[softly][to herself] Here we go again... [sighs] Is it Friday yet?' },
+  { label: '🦹 Villain', name: 'The Duke', desc: 'silky, theatrical villain with a rich low voice, amused and unhurried',
+    scene: '', text: '<slow>How delightful</slow>. You actually thought the plan was yours. [laughs]' },
+  { label: '👵 Nana', name: 'Nana', desc: 'warm elderly woman, soft and slightly shaky, a gentle storyteller',
+    scene: 'a quiet bedroom, the lights already dimmed', text: '[warm, gentle bedtime voice] Once upon a time, in a kingdom made entirely of cheese, there lived a very worried mouse.' },
+];
+const POINT_TAGS = ['laughs', 'chuckles', 'sighs', 'gasps', 'pause', 'clears throat', 'hesitates'];
+const SCOPED_TAGS = ['whisper', 'slow', 'emphasis'];
+const STYLES = ['with bright eager excitement', 'icy calm', 'heavy, deflated sadness', 'sarcastic and deadpan',
+  'breathless and terrified', 'furious, barely holding it together', 'warm, gentle bedtime voice'];
+const NOTES = ['softer and warmer', 'more excited', 'more amused than before; drop the hurt', 'slower, with more pauses',
+  'angrier', 'whispered', 'faster', 'more sarcastic'];
+// top_k 0 is no cap; settings saved before it became the default are not reused.
+const DEFAULT_SAMPLING = { temperature: 1.0, top_p: 0.9, top_k: 0 };
+const SAMPLING_KEY = 'sampling-v2';
+
+// ── Helpers ─────────────────────────────────────────────────────────────────
+const $ = id => document.getElementById(id);
+const saved = {
+  get(key, fallback) { try { const v = localStorage.getItem('drama3:' + key); return v == null ? fallback : JSON.parse(v); } catch { return fallback; } },
+  set(key, value) { try { localStorage.setItem('drama3:' + key, JSON.stringify(value)); } catch {} },
+};
+function el(tag, props = {}, text) {
+  const node = document.createElement(tag);
+  Object.assign(node, props);
+  if (text != null) node.textContent = text;
+  return node;
+}
+function showMsg(kind, text, action) {
+  const m = $('msg');
+  m.className = 'msg show ' + kind;
+  m.textContent = text;
+  if (action) m.append(el('a', { href: action.href, target: '_blank', rel: 'noopener' }, action.label));
+}
+function hideMsg() { $('msg').className = 'msg'; }
+function fmtTime(s) { s = Math.max(0, Math.round(s)); return Math.floor(s / 60) + ':' + String(s % 60).padStart(2, '0'); }
+
+// ── State ───────────────────────────────────────────────────────────────────
+const state = {
+  mode: saved.get('mode', 'random'),
+  voiceId: saved.get('voiceId', 'clone1'),
+  presets: [],
+  custom: [],
+  engineReady: false,
+  transcriber: 'waiting',
+  busy: false,
+  controller: null,
+  takes: [],
+  current: null,
+  sampling: { ...DEFAULT_SAMPLING, ...saved.get(SAMPLING_KEY, {}) },
+};
+
+// ── Theme ───────────────────────────────────────────────────────────────────
+function applyTheme(theme) { if (theme) document.documentElement.dataset.theme = theme; }
+function toggleTheme() {
+  const dark = document.documentElement.dataset.theme
+    ? document.documentElement.dataset.theme === 'dark'
+    : matchMedia('(prefers-color-scheme: dark)').matches;
+  const next = dark ? 'light' : 'dark';
+  applyTheme(next); saved.set('theme', next);
+}
+applyTheme(saved.get('theme', null));
+
+// ── Sheets ──────────────────────────────────────────────────────────────────
+function openSheet(id) { $(id).classList.add('open'); }
+function closeSheet(event) {
+  if (event && !event.target.classList.contains('sheet-bg')) return;
+  document.querySelectorAll('.sheet-bg.open').forEach(s => s.classList.remove('open'));
+}
+
+// ── Script + palette ────────────────────────────────────────────────────────
+function isMood(inner) {
+  const t = inner.trim().toLowerCase();
+  return STYLES.some(x => x.toLowerCase() === t) || t.split(/\s+/).length >= 3;
+}
+function esc(s) { return s.replace(/&/g, '&amp;').replace(/</g, '&lt;').replace(/>/g, '&gt;'); }
+function render() {
+  let h = esc($('script').value);
+  h = h.replace(/(^|\n)([ \t]*)(Direction:|Line:)/g, '$1$2<mark class="h-note">$3</mark>');
+  h = h.replace(/&lt;([A-Za-z][\w-]*)&gt;([\s\S]*?)&lt;\/\1&gt;/g, (_, name, body) =>
+    `<mark class="h-tag">&lt;${name}&gt;</mark><span class="h-span">${body}</span><mark class="h-tag">&lt;/${name}&gt;</mark>`);
+  h = h.replace(/\[([^\[\]\n]*)\]/g, (m, inner, offset, str) => {
+    const atStart = /^[ \t]*$/.test(str.slice(str.lastIndexOf('\n', offset - 1) + 1, offset));
+    return `<mark class="${atStart && isMood(inner) ? 'h-style' : 'h-point'}">[${inner}]</mark>`;
+  });
+  $('mirror').innerHTML = h + '\n\u200b';
+}
+function updateCount() {
+  render();
+  const n = $('script').value.length;
+  $('charCount').textContent = `${n} / ${state.maxChars || 600}`;
+  $('charCount').style.color = n > (state.maxChars || 600) ? 'var(--err)' : '';
+  saved.set('script', $('script').value);
+}
+function insertAtCursor(snippet) {
+  const ta = $('script');
+  const start = ta.selectionStart, end = ta.selectionEnd, v = ta.value;
+  const before = v.slice(0, start), after = v.slice(end);
+  const padL = before && !/\s$/.test(before) ? ' ' : '';
+  const padR = after && !/^\s/.test(after) ? ' ' : '';
+  ta.value = before + padL + snippet + padR + after;
+  const caret = (before + padL + snippet + padR).length;
+  ta.focus(); ta.setSelectionRange(caret, caret); updateCount();
+}
+function wrapSelection(tag) {
+  const ta = $('script');
+  const start = ta.selectionStart, end = ta.selectionEnd, v = ta.value;
+  const inner = v.slice(start, end) || 'these words';
+  ta.value = v.slice(0, start) + `<${tag}>${inner}</${tag}>` + v.slice(end);
+  const s = start + tag.length + 2;
+  ta.focus(); ta.setSelectionRange(s, s + inner.length); updateCount();
+}
+function setStyle(style) {
+  const ta = $('script');
+  const v = ta.value.replace(/^\s*\[[^\[\]\n]*\]\s*/, '');
+  ta.value = `[${style}] ` + v;
+  ta.focus(); updateCount();
+}
+function renderPalette() {
+  POINT_TAGS.forEach(t => $('pointTags').append(el('button', { type: 'button', className: 'chip tag', onclick: () => insertAtCursor(`[${t}]`) }, `[${t}]`)));
+  SCOPED_TAGS.forEach(t => $('scopedTags').append(el('button', { type: 'button', className: 'chip tag', title: 'Wraps the selected words', onclick: () => wrapSelection(t) }, `<${t}>…</${t}>`)));
+  STYLES.forEach(t => $('styleTags').append(el('button', { type: 'button', className: 'chip', title: 'Sets the delivery of the whole line', onclick: () => setStyle(t) }, t)));
+  EXAMPLES.forEach(x => $('examples').append(el('button', { type: 'button', className: 'chip dashed', onclick: () => {
+    $('script').value = x.text; updateCount();
+    if (x.note) $('noteIn').value = x.note;
+  } }, x.label)));
+  CAST.forEach(c => $('castChips').append(el('button', { type: 'button', className: 'chip', onclick: () => {
+    $('charName').value = c.name; $('charDesc').value = c.desc; $('charScene').value = c.scene;
+    $('script').value = c.text; updateCount(); saveDesign();
+  } }, c.label)));
+  NOTES.forEach(n => $('noteChips').append(el('button', { type: 'button', className: 'chip', onclick: () => { $('noteIn').value = n; retake(); } }, n)));
+}
+
+// ── Voice modes ─────────────────────────────────────────────────────────────
+function setMode(mode) {
+  state.mode = mode; saved.set('mode', mode);
+  document.querySelectorAll('.seg button').forEach(b => b.classList.toggle('on', b.dataset.mode === mode));
+  ['clone', 'design', 'random'].forEach(m => $('pane-' + m).classList.toggle('on', m === mode));
+}
+function allVoices() { return [...state.presets, ...state.custom]; }
+function voiceById(id) { return allVoices().find(v => v.id === id); }
+function renderVoices() {
+  const box = $('voiceChips');
+  box.innerHTML = '';
+  allVoices().forEach(v => {
+    box.append(el('button', {
+      type: 'button', className: 'chip' + (v.id === state.voiceId ? ' on' : ''),
+      title: v.transcript, onclick: () => selectVoice(v.id),
+    }, v.preset ? v.label : `★ ${v.label}`));
+  });
+  box.append(el('button', { type: 'button', className: 'chip dashed', onclick: () => $('fileIn').click() }, '⤒ Upload'));
+  const rec = el('span', { className: 'rec', id: 'recBox' });
+  rec.append(el('button', { type: 'button', className: 'chip dashed', id: 'recBtn', onclick: toggleRec }, '● Record'));
+  const meter = el('span', { className: 'meter' }); meter.append(el('i', { id: 'recLevel' }));
+  rec.append(meter);
+  box.append(rec);
+}
+function selectVoice(id) {
+  const v = voiceById(id);
+  if (!v) return;
+  state.voiceId = id; saved.set('voiceId', id);
+  $('transcript').value = v.transcript || '';
+  $('transcriptState').innerHTML = v.preset ? '' : (v.transcribed ? '<span class="badge">auto-transcribed</span>' : '<span class="badge">type it in</span>');
+  $('voicePreview').src = `api/voices/${encodeURIComponent(id)}.wav`;
+  renderVoices();
+}
+function saveDesign() { saved.set('design', { name: $('charName').value, desc: $('charDesc').value, scene: $('charScene').value }); }
+function voicePayload() {
+  if (state.mode === 'clone') return { mode: 'clone', id: state.voiceId, transcript: $('transcript').value };
+  if (state.mode === 'design') return { mode: 'design', name: $('charName').value || 'Narrator', description: $('charDesc').value, scene: $('charScene').value };
+  return { mode: 'random' };
+}
+function voiceLabel() {
+  if (state.mode === 'clone') return '🎙 ' + (voiceById(state.voiceId)?.label || 'Clone');
+  if (state.mode === 'design') return '✨ ' + ($('charName').value || 'Narrator');
+  return '🎲 Surprise voice';
+}
+
+// ── Adding voices ───────────────────────────────────────────────────────────
+$('fileIn').addEventListener('change', e => { const f = e.target.files[0]; e.target.value = ''; if (f) uploadVoice(f, f.name.replace(/\.[^.]+$/, '')); });
+async function uploadVoice(blob, label) {
+  hideMsg();
+  $('transcript').value = ''; $('transcript').placeholder = 'Transcribing…';
+  $('transcriptState').innerHTML = '<span class="badge">listening…</span>';
+  try {
+    const fd = new FormData();
+    fd.append('audio', blob, (label || 'voice') + '.wav');
+    const res = await fetch('api/voices', { method: 'POST', body: fd });
+    const d = await res.json().catch(() => ({}));
+    if (!res.ok) throw new Error(d.detail || 'Upload failed.');
+    addCustomVoice(d, label);
+    if (d.trimmed) showMsg('info', `Trimmed to ${d.duration.toFixed(1)} s at a pause, so the transcript matches.`);
+  } catch (e) {
+    showMsg('err', e.message);
+    selectVoice(state.voiceId);
+  } finally {
+    $('transcript').placeholder = 'Exact words spoken in the clip';
+  }
+}
+function addCustomVoice(d, label) {
+  state.custom = state.custom.filter(v => v.id !== d.id);
+  state.custom.push({ ...d, label: label || d.label });
+  setMode('clone');
+  selectVoice(d.id);
+}
+
+// ── Recording ───────────────────────────────────────────────────────────────
+const rec = { stream: null, ctx: null, proc: null, chunks: [], on: false, started: 0, raf: 0, peak: 0 };
+async function toggleRec() { if (rec.on) stopRec(); else startRec(); }
+async function startRec() {
+  if (!navigator.mediaDevices?.getUserMedia) { showMsg('err', 'Recording needs HTTPS and a microphone.'); return; }
+  hideMsg();
+  try {
+    rec.stream = await navigator.mediaDevices.getUserMedia({ audio: { channelCount: 1, echoCancellation: false, noiseSuppression: false, autoGainControl: false } });
+  } catch (e) { showMsg('err', 'Microphone blocked: ' + e.message); return; }
+  rec.ctx = new (window.AudioContext || window.webkitAudioContext)();
+  const src = rec.ctx.createMediaStreamSource(rec.stream);
+  rec.proc = rec.ctx.createScriptProcessor(4096, 1, 1);
+  const mute = rec.ctx.createGain(); mute.gain.value = 0;
+  rec.chunks = []; rec.on = true; rec.started = performance.now(); rec.peak = 0;
+  rec.proc.onaudioprocess = e => {
+    if (!rec.on) return;
+    const x = new Float32Array(e.inputBuffer.getChannelData(0));
+    let p = 0; for (let i = 0; i < x.length; i++) p = Math.max(p, Math.abs(x[i]));
+    rec.peak = Math.max(p, rec.peak * 0.85);
+    rec.chunks.push(x);
+  };
+  src.connect(rec.proc); rec.proc.connect(mute); mute.connect(rec.ctx.destination);
+  $('recBox').classList.add('live');
+  const tick = () => {
+    if (!rec.on) return;
+    const s = (performance.now() - rec.started) / 1000;
+    $('recBtn').textContent = `■ Stop ${fmtTime(s)}`;
+    $('recLevel').style.width = Math.min(100, rec.peak * 140) + '%';
+    if (s > 25) { stopRec(); return; }
+    rec.raf = requestAnimationFrame(tick);
+  };
+  tick();
+}
+function stopRec() {
+  rec.on = false; cancelAnimationFrame(rec.raf);
+  rec.stream?.getTracks().forEach(t => t.stop());
+  const sr = rec.ctx?.sampleRate || 48000;
+  rec.ctx?.close().catch(() => {});
+  $('recBox')?.classList.remove('live');
+  const total = rec.chunks.reduce((n, c) => n + c.length, 0);
+  renderVoices();
+  if (total < sr * 0.8) { showMsg('err', 'That was too short. Record at least a couple of seconds.'); return; }
+  const pcm = new Float32Array(total);
+  let off = 0; rec.chunks.forEach(c => { pcm.set(c, off); off += c.length; });
+  const n = state.custom.filter(v => v.label.startsWith('Mic')).length + 1;
+  uploadVoice(wavFromFloat(pcm, sr), `Mic ${n}`);
+}
+function wavFromFloat(pcm, sr) {
+  const buf = new ArrayBuffer(44 + pcm.length * 2), v = new DataView(buf);
+  const w = (o, s) => { for (let i = 0; i < s.length; i++) v.setUint8(o + i, s.charCodeAt(i)); };
+  w(0, 'RIFF'); v.setUint32(4, 36 + pcm.length * 2, true); w(8, 'WAVE'); w(12, 'fmt ');
+  v.setUint32(16, 16, true); v.setUint16(20, 1, true); v.setUint16(22, 1, true);
+  v.setUint32(24, sr, true); v.setUint32(28, sr * 2, true); v.setUint16(32, 2, true); v.setUint16(34, 16, true);
+  w(36, 'data'); v.setUint32(40, pcm.length * 2, true);
+  for (let i = 0; i < pcm.length; i++) { const s = Math.max(-1, Math.min(1, pcm[i])); v.setInt16(44 + i * 2, s < 0 ? s * 0x8000 : s * 0x7fff, true); }
+  return new Blob([buf], { type: 'audio/wav' });
+}
+function wavFromInt16(parts, sr) {
+  const bytes = parts.reduce((n, p) => n + p.byteLength, 0);
+  const head = new ArrayBuffer(44), v = new DataView(head);
+  const w = (o, s) => { for (let i = 0; i < s.length; i++) v.setUint8(o + i, s.charCodeAt(i)); };
+  w(0, 'RIFF'); v.setUint32(4, 36 + bytes, true); w(8, 'WAVE'); w(12, 'fmt ');
+  v.setUint32(16, 16, true); v.setUint16(20, 1, true); v.setUint16(22, 1, true);
+  v.setUint32(24, sr, true); v.setUint32(28, sr * 2, true); v.setUint16(32, 2, true); v.setUint16(34, 16, true);
+  w(36, 'data'); v.setUint32(40, bytes, true);
+  return new Blob([head, ...parts], { type: 'audio/wav' });
+}
+
+// ── Streaming playback ──────────────────────────────────────────────────────
+const player = {
+  ctx: null, node: null, analyser: null, queue: [], firstSoundAt: null, onFirstSound: null, raf: 0,
+  ensure(sr) {
+    if (this.ctx && this.ctx.sampleRate === sr) { this.ctx.resume(); return; }
+    this.ctx?.close().catch(() => {});
+    const AC = window.AudioContext || window.webkitAudioContext;
+    this.ctx = new AC({ sampleRate: sr });
+    this.analyser = this.ctx.createAnalyser();
+    this.analyser.fftSize = 1024; this.analyser.smoothingTimeConstant = 0.7;
+    this.node = this.ctx.createScriptProcessor(1024, 0, 1);
+    this.node.onaudioprocess = e => {
+      const out = e.outputBuffer.getChannelData(0);
+      let i = 0, wrote = false;
+      while (i < out.length && this.queue.length) {
+        const seg = this.queue[0];
+        const n = Math.min(out.length - i, seg.data.length - seg.pos);
+        out.set(seg.data.subarray(seg.pos, seg.pos + n), i);
+        seg.pos += n; i += n; wrote = true;
+        if (seg.pos >= seg.data.length) this.queue.shift();
+      }
+      if (i < out.length) out.fill(0, i);
+      if (wrote && this.firstSoundAt == null) { this.firstSoundAt = performance.now(); this.onFirstSound?.(); }
+    };
+    this.node.connect(this.analyser);
+    this.analyser.connect(this.ctx.destination);
+    this.ctx.resume();
+  },
+  reset() { this.queue = []; this.firstSoundAt = null; },
+  push(f32) { this.queue.push({ data: f32, pos: 0 }); },
+  buffered() { return this.queue.reduce((n, s) => n + s.data.length - s.pos, 0) / (this.ctx?.sampleRate || 44100); },
+};
+function drawViz() {
+  if (!$('stage').classList.contains('show')) { player.raf = requestAnimationFrame(drawViz); return; }
+  const canvas = $('viz');
+  const dpr = window.devicePixelRatio || 1;
+  const w = canvas.clientWidth, h = canvas.clientHeight;
+  if (canvas.width !== w * dpr) { canvas.width = w * dpr; canvas.height = h * dpr; }
+  const g = canvas.getContext('2d');
+  g.setTransform(dpr, 0, 0, dpr, 0, 0);
+  g.clearRect(0, 0, w, h);
+  const accent = getComputedStyle(document.documentElement).getPropertyValue('--accent').trim();
+  const faint = getComputedStyle(document.documentElement).getPropertyValue('--border').trim();
+  const bars = Math.max(24, Math.floor(w / 9));
+  const data = new Uint8Array(player.analyser ? player.analyser.frequencyBinCount : 0);
+  if (player.analyser) player.analyser.getByteFrequencyData(data);
+  const live = player.queue.length > 0;
+  for (let b = 0; b < bars; b++) {
+    const idx = Math.floor(Math.pow(b / bars, 1.4) * data.length * 0.35);
+    const v = data.length ? data[idx] / 255 : 0;
+    const bh = Math.max(2, v * (h - 10));
+    g.fillStyle = live && v > 0.02 ? accent : faint;
+    const x = b * (w / bars) + 2;
+    g.beginPath();
+    g.roundRect ? g.roundRect(x, (h - bh) / 2, w / bars - 4, bh, 2) : g.rect(x, (h - bh) / 2, w / bars - 4, bh);
+    g.fill();
+  }
+  player.raf = requestAnimationFrame(drawViz);
+}
+
+// ── Perform ─────────────────────────────────────────────────────────────────
+function setPerform(label, live) {
+  $('performLabel').textContent = label;
+  $('performBtn').classList.toggle('live', !!live);
+  $('performBtn').disabled = state.busy || !state.engineReady;
+  $('stopBtn').classList.toggle('show', state.busy);
+  $('retakeBtn').disabled = state.busy || !state.engineReady || !state.current?.id;
+  document.querySelectorAll('#noteChips .chip').forEach(c => c.disabled = state.busy || !state.current?.id);
+}
+function idleLabel() { return state.engineReady ? 'Perform' : 'Warming up…'; }
+function b64ToInt16(b64, carry) {
+  const bin = atob(b64);
+  const bytes = new Uint8Array(carry.length + bin.length);
+  bytes.set(carry);
+  for (let i = 0; i < bin.length; i++) bytes[carry.length + i] = bin.charCodeAt(i);
+  const even = bytes.length - (bytes.length % 2);
+  return { pcm: new Int16Array(bytes.buffer, 0, even / 2), carry: bytes.slice(even) };
+}
+
+async function perform(options = {}) {
+  if (state.busy || !state.engineReady) return;
+  const text = $('script').value.trim();
+  if (!text) { showMsg('err', 'Write a line for the actor first.'); $('script').focus(); return; }
+  if (!options.retake) {
+    if (state.mode === 'clone' && !$('transcript').value.trim()) { showMsg('err', 'Add what the reference clip says, so Drama 3 can line the voice up with the words.'); $('transcript').focus(); return; }
+    if (state.mode === 'design' && !$('charDesc').value.trim()) { showMsg('err', 'Describe the voice you want to cast.'); $('charDesc').focus(); return; }
+  }
+  hideMsg();
+  const body = { text, voice: voicePayload(), ...state.sampling };
+  const parent = options.retake ? state.current : null;
+  if (parent) { body.retake_of = parent.id; body.note = options.note; }
+
+  player.ensure(44100);          // inside the click, so autoplay rules allow sound
+  player.reset();
+  state.busy = true;
+  const controller = new AbortController();
+  state.controller = controller;
+  const take = {
+    id: null, n: state.takes.length + 1, voice: parent ? parent.voice : voiceLabel(),
+    mode: parent ? parent.mode : state.mode, note: parent ? options.note : null,
+    character: parent ? parent.character : (state.mode === 'design' ? ($('charName').value || 'Narrator') : null),
+    text, parts: [], sr: 44100, metrics: {}, blob: null,
+  };
+  const t0 = performance.now();
+  let firstChunkAt = null, carry = new Uint8Array(0);
+  player.onFirstSound = () => { take.metrics.client_ms = Math.round(player.firstSoundAt - t0); renderMetrics(take, true); };
+  showStage(take);
+  setPerform('Connecting…', true);
+
+  try {
+    const res = await fetch('api/generate', {
+      method: 'POST', headers: { 'Content-Type': 'application/json' },
+      body: JSON.stringify(body), signal: controller.signal,
+    });
+    if (res.status === 401) { location.reload(); return; }
+    if (!res.ok) {
+      const e = await res.json().catch(() => ({}));
+      throw Object.assign(new Error(typeof e.detail === 'string' ? e.detail : `Request failed (${res.status})`), { status: res.status });
+    }
+    const reader = res.body.getReader();
+    const decoder = new TextDecoder();
+    let buf = '';
+    for (;;) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      buf += decoder.decode(value, { stream: true });
+      let cut;
+      while ((cut = buf.indexOf('\n\n')) >= 0) {
+        const block = buf.slice(0, cut); buf = buf.slice(cut + 2);
+        for (const line of block.split('\n')) {
+          if (!line.startsWith('data: ')) continue;
+          const d = JSON.parse(line.slice(6));
+          if (d.type === 'queued') {
+            setPerform(d.position <= 1 ? 'Next in line…' : `In line · ${d.position - 1} ahead`, true);
+          } else if (d.type === 'retrying') {
+            setPerform('Trying once more…', true);
+          } else if (d.type === 'start') {
+            take.id = d.take_id; take.sr = d.sample_rate;
+            if (d.usage) renderUsage(d.usage);
+            player.ensure(d.sample_rate);
+            setPerform('Performing…', true);
+          } else if (d.type === 'audio') {
+            const out = b64ToInt16(d.pcm, carry); carry = out.carry;
+            if (firstChunkAt == null) firstChunkAt = performance.now();
+            if (d.ttfa_ms != null) take.metrics.ttfa_ms = d.ttfa_ms;
+            const f32 = new Float32Array(out.pcm.length);
+            for (let i = 0; i < out.pcm.length; i++) f32[i] = out.pcm[i] / 32768;
+            player.push(f32);
+            take.parts.push(new Int16Array(out.pcm));
+            take.metrics.audio_s = take.parts.reduce((n, p) => n + p.length, 0) / take.sr;
+            renderMetrics(take, true);
+          } else if (d.type === 'done') {
+            Object.assign(take.metrics, d);
+            take.turn = d.turn_text;
+            if (d.usage) renderUsage(d.usage);
+            if (d.warning) showMsg('info', d.warning);
+          } else if (d.type === 'error') {
+            throw new Error(d.message);
+          }
+        }
+      }
+    }
+    if (!take.metrics.total_ms) throw new Error('The take was cut off. Try again.');
+    finishTake(take);
+  } catch (e) {
+    if (e.name === 'AbortError') {
+      showMsg('info', 'Cut! Take stopped.');
+      if (take.parts.length) finishTake(take, true);
+    } else {
+      const signIn = e.status === 429 && state.user && !state.user.signed_in
+        ? { label: 'Sign in with Hugging Face', href: '/oauth/huggingface/login' } : null;
+      showMsg('err', e.message || 'Something went wrong.', signIn);
+      if (!state.current) $('stage').classList.remove('show'); else showStage(state.current);
+    }
+  } finally {
+    state.busy = false; state.controller = null;
+    setPerform(idleLabel(), false);
+  }
+}
+function stopTake() { state.controller?.abort(); player.reset(); }
+function retake() {
+  const note = $('noteIn').value.trim();
+  if (!state.current?.id) return;
+  if (!note) { showMsg('err', 'Type a note for the actor, or tap one of the suggestions.'); $('noteIn').focus(); return; }
+  perform({ retake: true, note });
+}
+
+// ── Takes ───────────────────────────────────────────────────────────────────
+function showStage(take) {
+  $('stage').classList.add('show');
+  $('takeTitle').textContent = `Take ${take.n}` + (take.note ? ' · retake' : '');
+  $('takeVoice').textContent = take.voice;
+  renderMetrics(take, !take.blob);
+  $('takeAudio').hidden = !take.blob;
+  if (take.blob) { $('takeAudio').src = take.url; }
+  $('downloadBtn').disabled = !take.blob;
+  $('keepBtn').disabled = !take.blob || !take.id;
+  $('keepBtn').style.display = take.mode === 'clone' ? 'none' : '';
+  if (!take.blob) $('stage').scrollIntoView({ behavior: 'smooth', block: 'nearest' });
+}
+function finishTake(take, partial) {
+  take.blob = wavFromInt16(take.parts, take.sr);
+  take.url = URL.createObjectURL(take.blob);
+  take.partial = !!partial;
+  state.takes.unshift(take);
+  if (take.id && !partial) state.current = take;
+  showStage(take);
+  renderSlate();
+}
+function renderMetrics(take, live) {
+  const m = take.metrics, parts = [];
+  const add = (k, v, hot) => parts.push(`<span class="${hot ? 'hot' : ''}">${k} <b>${v}</b></span>`);
+  if (m.ttfa_ms != null) add('first audio', `${m.ttfa_ms} ms`, true);
+  if (m.client_ms != null) add('you heard it', `${m.client_ms} ms`);
+  if (m.speed != null) add('speed', `${m.speed}× real-time`, true);
+  if (m.audio_s != null) add('length', `${m.audio_s.toFixed(1)} s`);
+  if (m.queue_ms) add('queue', `${(m.queue_ms / 1000).toFixed(1)} s`);
+  if (live) add('buffer', `${player.buffered().toFixed(1)} s`);
+  if (m.trimmed_s >= 0.5) add('silence trimmed', `${m.trimmed_s.toFixed(1)} s`);
+  if (m.finish_reason === 'length') add('note', 'hit the length cap');
+  if (m.finish_reason === 'stall') add('note', 'stalled, cut short');
+  if (take.note) parts.push(`<span>note <b>${escapeHtml(take.note)}</b></span>`);
+  $('metrics').innerHTML = parts.join('');
+}
+function escapeHtml(s) { return s.replace(/[&<>"']/g, c => ({ '&': '&amp;', '<': '&lt;', '>': '&gt;', '"': '&quot;', "'": '&#39;' }[c])); }
+function renderSlate() {
+  const box = $('slate'); box.innerHTML = '';
+  state.takes.forEach(take => {
+    const item = el('div', { className: 'slate-item' + (take === state.current ? ' current' : '') });
+    item.append(el('span', { className: 'slate-num' }, `#${take.n}`));
+    const body = el('div', { className: 'slate-body' });
+    body.append(el('div', { className: 'slate-title' }, take.note ? `“${take.note}”` : take.text.replace(/\s+/g, ' ')));
+    body.append(el('div', { className: 'slate-sub' }, `${take.voice} · ${(take.metrics.audio_s || 0).toFixed(1)} s` +
+      (take.metrics.ttfa_ms != null ? ` · ${take.metrics.ttfa_ms} ms to first audio` : '') + (take.partial ? ' · stopped' : '')));
+    item.append(body);
+    const actions = el('div', { className: 'slate-actions' });
+    actions.append(el('button', { type: 'button', className: 'mini', onclick: () => {
+      if (take.id && !take.partial) state.current = take;
+      showStage(take); renderSlate(); $('takeAudio').play().catch(() => {});
+    } }, take === state.current ? '▶ Play' : '▶ Open'));
+    item.append(actions);
+    box.append(item);
+  });
+}
+function downloadTake() {
+  const take = state.takes.find(t => t.url === $('takeAudio').src) || state.current;
+  if (!take?.blob) return;
+  const a = el('a', { href: take.url, download: `drama3-take-${take.n}.wav` });
+  document.body.append(a); a.click(); a.remove();
+}
+async function keepVoice() {
+  const take = state.current;
+  if (!take?.id) return;
+  $('keepBtn').disabled = true; $('keepBtn').textContent = 'Keeping…';
+  try {
+    const label = take.character || `Found ${state.custom.length + 1}`;
+    const res = await fetch('api/voices/from-take', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ take_id: take.id, label }) });
+    const d = await res.json().catch(() => ({}));
+    if (!res.ok) throw new Error(d.detail || 'Could not keep that voice.');
+    addCustomVoice(d, label);
+    showMsg('info', `Saved “${label}” as a clone voice. New lines will keep this voice.`);
+  } catch (e) { showMsg('err', e.message); }
+  finally { $('keepBtn').textContent = 'Keep this voice'; $('keepBtn').disabled = false; }
+}
+
+// ── Status ──────────────────────────────────────────────────────────────────
+async function poll() {
+  let delay = 2500;
+  try {
+    const res = await fetch('api/status', { cache: 'no-store' });
+    if (res.status === 401) { location.reload(); return; }
+    const d = await res.json();
+    applyStatus(d);
+    delay = state.engineReady ? 8000 : 2500;
+  } catch {
+    $('statusDot').className = 'dot err'; $('statusText').textContent = 'offline';
+  }
+  setTimeout(poll, delay);
+}
+function applyStatus(d) {
+  renderUser(d.user, d.usage);
+  const e = d.engine;
+  state.maxChars = d.limits?.max_text_chars;
+  if (!state.presets.length && d.voices?.length) {
+    state.presets = d.voices;
+    if (!voiceById(state.voiceId)) state.voiceId = d.voices[0].id;
+    selectVoice(state.voiceId);
+  }
+  state.transcriber = d.transcriber;
+  const wasReady = state.engineReady;
+  state.engineReady = e.state === 'ready';
+  const boot = $('boot');
+  boot.classList.toggle('show', !state.engineReady);
+  boot.classList.toggle('error', e.state === 'error');
+  if (!state.engineReady) {
+    $('bootTitle').textContent = e.state === 'error' ? 'The stage is dark' : e.state === 'restarting' ? 'Restarting the engine…' : 'Setting the stage…';
+    $('bootPhase').textContent = e.state === 'error' ? e.phase : `${e.phase} · ${fmtTime(e.elapsed_s)}. The page works meanwhile: write a script and pick a voice.`;
+    $('bootLog').textContent = e.log || '';
+  }
+  const gpu = d.gpu?.name || 'GPU';
+  const dot = $('statusDot'), txt = $('statusText');
+  if (state.engineReady) {
+    const busy = d.load.active + d.load.waiting;
+    dot.className = 'dot ready';
+    const bits = [`<b>${escapeHtml(gpu)}</b>`];
+    bits.push(busy ? `${busy} performing` : 'ready');
+    if (d.stats.ttfa_ms_p50 != null) bits.push(`${d.stats.ttfa_ms_p50} ms to first audio`);
+    txt.innerHTML = bits.join(' · ');
+  } else {
+    dot.className = 'dot ' + (e.state === 'error' ? 'err' : 'busy');
+    txt.innerHTML = `<b>${escapeHtml(gpu)}</b> · ${e.state === 'error' ? 'down' : 'warming up'}`;
+  }
+  if (wasReady !== state.engineReady && !state.busy) setPerform(idleLabel(), false);
+  updateCount();
+}
+
+// ── Account ─────────────────────────────────────────────────────────────────
+function renderUser(user, usage) {
+  state.user = user && user.username !== 'local' ? user : null;
+  $('account').hidden = !state.user || !state.user.signed_in;
+  $('signIn').hidden = !state.user || state.user.signed_in;
+  renderUsage(usage);
+}
+function renderUsage(usage) {
+  if (!usage || !state.user) return;
+  if (!state.user.signed_in) {
+    const left = Math.max(0, usage.remaining ?? 0);
+    $('signIn').textContent = `Sign in · ${left} guest take${left === 1 ? '' : 's'} left`;
+    $('signIn').title = `Signing in with Hugging Face adds ${usage.free_after_sign_in} takes a day; PRO is unlimited.`;
+    return;
+  }
+  const left = Math.max(0, usage.remaining ?? 0);
+  const badge = usage.unlimited ? (usage.tier === 'pro' ? 'PRO' : 'team') : `${left} left`;
+  $('userPill').textContent = `${state.user.username} · ${badge}`;
+  $('accountSummary').textContent = usage.unlimited
+    ? `${badge} · unlimited takes`
+    : `Free · ${left} of ${usage.limit} takes left today (resets at midnight UTC)`;
+}
+function toggleAccount(event) {
+  event.stopPropagation();
+  const menu = $('accountMenu');
+  menu.hidden = !menu.hidden;
+  $('userPill').setAttribute('aria-expanded', String(!menu.hidden));
+}
+document.addEventListener('click', e => { if (!$('account').contains(e.target)) $('accountMenu').hidden = true; });
+
+// ── Settings ────────────────────────────────────────────────────────────────
+function bindSlider(id, key, fmt) {
+  const input = $(id), out = $(id + 'V');
+  const sync = () => { out.textContent = fmt(state.sampling[key]); input.value = state.sampling[key]; };
+  input.addEventListener('input', () => { state.sampling[key] = Number(input.value); saved.set(SAMPLING_KEY, state.sampling); sync(); });
+  sync();
+  return sync;
+}
+const syncs = [
+  bindSlider('temp', 'temperature', v => v.toFixed(2)),
+  bindSlider('topp', 'top_p', v => v.toFixed(2)),
+  bindSlider('topk', 'top_k', v => v ? String(v) : 'off'),
+];
+function resetSampling() { state.sampling = { ...DEFAULT_SAMPLING }; saved.set(SAMPLING_KEY, state.sampling); syncs.forEach(s => s()); }
+
+// ── Init ────────────────────────────────────────────────────────────────────
+renderPalette();
+const design = saved.get('design', null) || CAST[0];
+$('charName').value = design.name; $('charDesc').value = design.desc; $('charScene').value = design.scene || '';
+['charName', 'charDesc', 'charScene'].forEach(id => $(id).addEventListener('input', saveDesign));
+$('script').value = saved.get('script', EXAMPLES[0].text);
+$('script').addEventListener('input', updateCount);
+updateCount();
+setMode(['clone', 'design', 'random'].includes(state.mode) ? state.mode : 'random');
+renderVoices();
+setPerform(idleLabel(), false);
+drawViz();
+poll();
+document.addEventListener('keydown', e => {
+  if ((e.metaKey || e.ctrlKey) && e.key === 'Enter') { e.preventDefault(); if (document.activeElement === $('noteIn')) retake(); else perform(); }
+  if (e.key === 'Escape') { closeSheet(); $('accountMenu').hidden = true; }
+});
+$('noteIn').addEventListener('keydown', e => { if (e.key === 'Enter' && !e.metaKey && !e.ctrlKey) { e.preventDefault(); retake(); } });
+</script>
+</body>
+</html>

requirements-preview.txt

@@ -0,0 +1,8 @@
+# Shared UI dependencies. CPU deployment adds requirements-ui.txt; GPU uses Dockerfile.engine.
+gradio[oauth]==6.28.0
+huggingface-hub[oauth]>=1.5.0,<2
+numpy
+soundfile
+librosa
+uvicorn
+python-multipart

requirements-ui.txt

@@ -0,0 +1,7 @@
+-r requirements-preview.txt
+faster-whisper==1.2.1
+av
+httpx[http2]
+fastapi
+itsdangerous
+authlib

requirements.txt

@@ -1,15 +0,0 @@
-gradio==6.30.0
-spaces
-torch==2.13.0
-torchvision==0.28.0
-transformers>=5.0
-huggingface_hub
-safetensors
-tokenizers
-loguru
-tqdm
-einops
-numpy
-librosa
-soundfile
-httpx

scripts/omni_missing_deps.py

@@ -0,0 +1,63 @@
+"""Pins of SGLang Omni's pyproject.toml that this environment does not satisfy.
+
+The base image moves slower than SGLang Omni's main branch, so its CI installs these
+before testing a commit (.github/scripts/omni_missing_dependencies.py there); the
+Dockerfile does the same for the pinned commit.
+
+    python3 omni_missing_deps.py pyproject.toml              # requirements to install
+    python3 omni_missing_deps.py --overrides pyproject.toml  # uv overrides, for --no-deps
+    python3 omni_missing_deps.py --report pyproject.toml     # what is still unsatisfied
+"""
+
+import argparse
+import importlib.metadata
+import tomllib
+
+from packaging.requirements import Requirement
+
+EXTRAS = ("fishaudio",)
+
+
+def unsatisfied(specs: list[str]) -> list[str]:
+    missing = []
+    for spec in specs:
+        requirement = Requirement(spec)
+        if requirement.marker and not requirement.marker.evaluate():
+            continue
+        try:
+            installed = importlib.metadata.version(requirement.name)
+        except importlib.metadata.PackageNotFoundError:
+            missing.append(spec)
+            continue
+        if not requirement.specifier.contains(installed, prereleases=True):
+            missing.append(spec)
+    return missing
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("pyproject")
+    parser.add_argument("--overrides", action="store_true")
+    parser.add_argument("--report", action="store_true")
+    args = parser.parse_args()
+    with open(args.pyproject, "rb") as f:
+        config = tomllib.load(f)
+    overrides = config.get("tool", {}).get("uv", {}).get("override-dependencies", [])
+    if args.overrides:
+        print("\n".join(overrides))
+        return
+    project = config["project"]
+    specs = list(project["dependencies"])
+    for extra in EXTRAS:
+        specs += project["optional-dependencies"][extra]
+    missing = unsatisfied(specs)
+    if args.report:
+        missing += unsatisfied(overrides)
+        print("All SGLang Omni pins satisfied." if not missing else
+              "Still unsatisfied (check the engine log if it fails to start):\n  " + "\n  ".join(missing))
+        return
+    print("\n".join(missing))
+
+
+if __name__ == "__main__":
+    main()