Compare commits
47
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b22474c0a5
|
||
|
|
0898d13718
|
||
|
|
9740812517
|
||
|
|
1065ffa12e
|
||
|
|
f9613c0a11
|
||
|
|
fe5a1a826f
|
||
|
|
de2378b8b3
|
||
|
|
a4fbe929ea
|
||
|
|
1bd8991053
|
||
|
|
3b73bd2728
|
||
|
|
c545e50d41
|
||
|
|
d88e8fbaa0
|
||
|
|
095f2f3a30
|
||
|
|
e568cf9d37
|
||
|
|
a25c20e8c8
|
||
|
|
8bd371824d
|
||
|
|
f40406e49b
|
||
|
|
6535efc51f
|
||
|
|
231c4b3991
|
||
|
|
03b3e3e701
|
||
|
|
1f96bf1d16
|
||
|
|
4431b19461
|
||
|
|
01a3ddaefa
|
||
|
|
1a1129f3d1
|
||
|
|
4904160431
|
||
|
|
c7fc3308b0
|
||
|
|
a1592cd38e
|
||
|
|
6fd81539e5
|
||
|
|
6c45663f3c
|
||
|
|
6498c27a4b
|
||
|
|
330fa08e91
|
||
|
|
8af0baba42 | ||
|
|
8e99cedadf
|
||
|
|
31cc853445
|
||
|
|
46f0164105
|
||
|
|
dbde8f061b
|
||
|
|
63dcc4f5e9 | ||
|
|
7596e7eafe
|
||
|
|
b922a5a29e
|
||
|
|
367ba7e4a7
|
||
|
|
9b508bc380
|
||
|
|
e6ff5a5962 | ||
|
|
def7ab9147 | ||
|
|
70bd8b599d
|
||
|
|
15a9787be2
|
||
|
|
cdbc853ff1
|
||
|
|
fa7af0335a
|
@@ -52,6 +52,11 @@ concurrency:
|
||||
env:
|
||||
CARGO_INCREMENTAL: "0"
|
||||
CARGO_TERM_COLOR: "always"
|
||||
# Every RPM here is built for Fedora 43, but the runners build on their
|
||||
# host, which may run a newer release. Each Stage step fails a binary that
|
||||
# needs a newer glibc than Fedora 43 ships, rather than letting dnf skip
|
||||
# the upgrade at deploy.
|
||||
GLIBC_FLOOR: "2.42"
|
||||
|
||||
jobs:
|
||||
prepare:
|
||||
@@ -315,6 +320,7 @@ jobs:
|
||||
run: |
|
||||
mkdir --parents artifacts
|
||||
cp target/release/cortex artifacts/cortex
|
||||
script/check-glibc-floor.sh "$GLIBC_FLOOR" artifacts/cortex
|
||||
./artifacts/cortex --version || true
|
||||
|
||||
- uses: actions/upload-artifact@v3
|
||||
@@ -354,6 +360,7 @@ jobs:
|
||||
run: |
|
||||
mkdir --parents artifacts
|
||||
cp target/release/helexa-bench artifacts/helexa-bench
|
||||
script/check-glibc-floor.sh "$GLIBC_FLOOR" artifacts/helexa-bench
|
||||
./artifacts/helexa-bench --version || true
|
||||
|
||||
- uses: actions/upload-artifact@v3
|
||||
@@ -393,6 +400,7 @@ jobs:
|
||||
run: |
|
||||
mkdir --parents artifacts
|
||||
cp target/release/helexa-upstream artifacts/helexa-upstream
|
||||
script/check-glibc-floor.sh "$GLIBC_FLOOR" artifacts/helexa-upstream
|
||||
./artifacts/helexa-upstream --version || true
|
||||
|
||||
- uses: actions/upload-artifact@v3
|
||||
@@ -432,6 +440,7 @@ jobs:
|
||||
run: |
|
||||
mkdir --parents artifacts
|
||||
cp target/release/helexa-angels artifacts/helexa-angels
|
||||
script/check-glibc-floor.sh "$GLIBC_FLOOR" artifacts/helexa-angels
|
||||
./artifacts/helexa-angels --version || true
|
||||
|
||||
- uses: actions/upload-artifact@v3
|
||||
@@ -467,6 +476,7 @@ jobs:
|
||||
run: |
|
||||
mkdir --parents artifacts
|
||||
cp target/release/helexa-router artifacts/helexa-router
|
||||
script/check-glibc-floor.sh "$GLIBC_FLOOR" artifacts/helexa-router
|
||||
./artifacts/helexa-router --version || true
|
||||
|
||||
- uses: actions/upload-artifact@v3
|
||||
@@ -502,6 +512,7 @@ jobs:
|
||||
run: |
|
||||
mkdir --parents artifacts
|
||||
cp target/release/helexa-tools artifacts/helexa-tools
|
||||
script/check-glibc-floor.sh "$GLIBC_FLOOR" artifacts/helexa-tools
|
||||
./artifacts/helexa-tools --version || true
|
||||
|
||||
- uses: actions/upload-artifact@v3
|
||||
@@ -619,6 +630,7 @@ jobs:
|
||||
run: |
|
||||
mkdir --parents artifacts
|
||||
cp target/release/neuron artifacts/neuron-${{ matrix.flavour }}
|
||||
script/check-glibc-floor.sh "$GLIBC_FLOOR" "artifacts/neuron-${{ matrix.flavour }}"
|
||||
file "artifacts/neuron-${{ matrix.flavour }}"
|
||||
|
||||
- uses: actions/upload-artifact@v3
|
||||
|
||||
@@ -249,6 +249,13 @@ jobs:
|
||||
runs-on: rust
|
||||
needs: changes
|
||||
if: needs.changes.outputs.rust == 'true'
|
||||
# Line tables only: full debuginfo for the whole graph made the neuron
|
||||
# test binary's compile and link the memory peak of the job, and on the
|
||||
# 15 GB rust runners a Test sharing the host with a Clippy was killed
|
||||
# mid-compile with no log (2026-10-17). Backtraces keep file:line.
|
||||
env:
|
||||
CARGO_PROFILE_DEV_DEBUG: line-tables-only
|
||||
CARGO_PROFILE_TEST_DEBUG: line-tables-only
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
# See script/ci-cargo-escalate.sh for the escalation rationale.
|
||||
|
||||
+55
-10
@@ -183,6 +183,13 @@ jobs:
|
||||
# enable --now: start the service AND enable it for boot so the
|
||||
# fleet self-heals after a host reboot.
|
||||
sudo /usr/bin/systemctl enable --now cortex.service
|
||||
# dnf can exit 0 without installing anything (e.g. an unmet glibc
|
||||
# requirement), so compare what is now installed with what was published.
|
||||
now=$(rpm -q --qf '%{VERSION}-%{RELEASE}' "${pkg}" 2>/dev/null || true)
|
||||
if [ -n "${latest}" ] && [ "${now}" != "${latest}" ]; then
|
||||
echo "::error::${pkg}: ${now:-nothing} is installed, not the published ${latest}"
|
||||
exit 1
|
||||
fi
|
||||
DEPLOY
|
||||
|
||||
# Wait for the service to either come up or wedge, then capture
|
||||
@@ -244,14 +251,12 @@ jobs:
|
||||
flavour: blackwell
|
||||
load_timeout: 900
|
||||
max_prompt_tokens: 131072
|
||||
# benjy is out of the matrix while its boot NVMe is replaced
|
||||
# (dead 2026-09-29): every deploy failed at SSH init against a
|
||||
# host that is not there. Restore these four lines when it is
|
||||
# back on the mesh.
|
||||
# - host: benjy.hanzalova.internal
|
||||
# flavour: ada
|
||||
# load_timeout: 300
|
||||
# max_prompt_tokens: 16384
|
||||
# benjy: the everyday class (Qwen3-8B) plus the cold-swap
|
||||
# targets that take turns with it -- Z-Image and ACE-Step music.
|
||||
- host: benjy.hanzalova.internal
|
||||
flavour: ada
|
||||
load_timeout: 300
|
||||
max_prompt_tokens: 16384
|
||||
- host: quadbrat.hanzalova.internal
|
||||
flavour: ampere
|
||||
load_timeout: 300
|
||||
@@ -464,6 +469,13 @@ jobs:
|
||||
# enable --now: start the service AND enable it for boot so the
|
||||
# fleet self-heals after a host reboot.
|
||||
sudo /usr/bin/systemctl enable --now neuron.service
|
||||
# dnf can exit 0 without installing anything (e.g. an unmet glibc
|
||||
# requirement), so compare what is now installed with what was published.
|
||||
now=$(rpm -q --qf '%{VERSION}-%{RELEASE}' "${pkg}" 2>/dev/null || true)
|
||||
if [ -n "${latest}" ] && [ "${now}" != "${latest}" ]; then
|
||||
echo "::error::${pkg}: ${now:-nothing} is installed, not the published ${latest}"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# ── Post-deploy validation ────────────────────────────────
|
||||
# A deploy only goes green if the neuron (a) finishes loading
|
||||
@@ -508,10 +520,28 @@ jobs:
|
||||
sys.stderr.write(msg + chr(10))
|
||||
sys.exit(1)
|
||||
completed = a.get("completed", [])
|
||||
print(completed[0] if completed else "")
|
||||
print(" ".join(completed))
|
||||
')
|
||||
# Probe the first default that answers chat. A decision model
|
||||
# (Laya on quadbrat) or an image or audio model 422s a chat
|
||||
# request, which made every quadbrat deploy report failure.
|
||||
chat_models=$(curl -s --max-time 10 http://localhost:13131/models | python3 -c '
|
||||
import json, sys
|
||||
for m in json.load(sys.stdin):
|
||||
caps = m.get("capabilities", [])
|
||||
if "text" in caps and not {"decision", "image", "audio"} & set(caps):
|
||||
print(m["id"])
|
||||
')
|
||||
probe_model=""
|
||||
for candidate in ${model}; do
|
||||
if printf '%s\n' "${chat_models}" | grep -qxF "${candidate}"; then
|
||||
probe_model="${candidate}"
|
||||
break
|
||||
fi
|
||||
done
|
||||
model="${probe_model}"
|
||||
if [ -z "${model}" ]; then
|
||||
echo "no default models configured — skipping LLM probe"
|
||||
echo "no chat-serving default model — skipping LLM probe"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
@@ -842,7 +872,15 @@ jobs:
|
||||
fi
|
||||
sudo /usr/bin/systemctl daemon-reload
|
||||
sudo /usr/bin/systemctl enable --now "${pkg}.service"
|
||||
# dnf can exit 0 without installing anything (e.g. an unmet glibc
|
||||
# requirement), so compare what is now installed with what was published.
|
||||
now=$(rpm -q --qf '%{VERSION}-%{RELEASE}' "${pkg}" 2>/dev/null || true)
|
||||
if [ -n "${latest}" ] && [ "${now}" != "${latest}" ]; then
|
||||
echo "::error::${pkg}: ${now:-nothing} is installed, not the published ${latest}"
|
||||
failed=1
|
||||
fi
|
||||
done
|
||||
[ "${failed:-0}" -eq 0 ]
|
||||
DEPLOY
|
||||
|
||||
# Router config (#287). Tracked because it holds no secrets, and
|
||||
@@ -1071,6 +1109,13 @@ jobs:
|
||||
# enable --now: start the service AND enable it for boot so the
|
||||
# bench resumes collecting after a host reboot.
|
||||
sudo /usr/bin/systemctl enable --now helexa-bench.service
|
||||
# dnf can exit 0 without installing anything (e.g. an unmet glibc
|
||||
# requirement), so compare what is now installed with what was published.
|
||||
now=$(rpm -q --qf '%{VERSION}-%{RELEASE}' "${pkg}" 2>/dev/null || true)
|
||||
if [ -n "${latest}" ] && [ "${now}" != "${latest}" ]; then
|
||||
echo "::error::${pkg}: ${now:-nothing} is installed, not the published ${latest}"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# ── Post-deploy validation ────────────────────────────────
|
||||
# The bench serves a read-only API on :13132 alongside the
|
||||
|
||||
@@ -1242,3 +1242,130 @@ retries, a bench `decision:` scenario, a dashboard panel,
|
||||
micro-batching); #342 (serving fine-tuned checkpoints); #352 (Z-Image on
|
||||
the ampere flash build, unrelated to Laya but enabled by the same
|
||||
flavour change).
|
||||
|
||||
## 2026-10-17 addendum: music generation (epics #378, #379)
|
||||
|
||||
Music is the fourth modality. Two engines serve candle-natively on
|
||||
benjy's RTX 4090 through the device-worker discipline (`AudioHandle`
|
||||
slab; `LoadAudio` / `DropAudio` / `GenerateAudio` jobs; CPU-side sample
|
||||
replies, FLAC encoded async-side):
|
||||
|
||||
- **ACE-Step 1.5** (MIT): turbo DiT + Qwen3-Embedding text encoder +
|
||||
48 kHz stereo VAE, with an optional 1.7B planner LM. Vendored from
|
||||
candle PR #3451 into `harness/arch/acestep/`. 10–600 s.
|
||||
- **HeartMuLa oss-3B** (Apache-2.0; Llama 3.2 tokenizer): an
|
||||
autoregressive LM (bf16) plus HeartCodec (f32), from two sibling repos
|
||||
(`HeartCodec-oss-20260123`, `HeartMuLaGen`). Detected by
|
||||
`model_type: heartmula`. 10–240 s.
|
||||
|
||||
**Surface: async jobs only.** `/v1/audio/generations`: POST submits and
|
||||
returns 202 with a job; GET lists; `/{id}` polls; `/{id}/content`
|
||||
fetches the audio; DELETE cancels. **cortex owns the jobs**, persisting
|
||||
them under `[audio_jobs] dir` and re-queueing unfinished ones on restart.
|
||||
neuron renders synchronously and returns the FLAC plus an
|
||||
`x-helexa-audio-result` header. The router dispatches POST by capacity
|
||||
and fans GET/DELETE out across cortexes until one answers non-404. Wire
|
||||
types are in `cortex-core/src/audio.rs`; limits (`[models.audio]`) and
|
||||
model-scoped licence notices (`[[models.notices]]`) come from the
|
||||
catalogue and are advertised in `/v1/models`.
|
||||
|
||||
**Formats.** Audio is rendered straight to 16-bit FLAC (no WAV step) and
|
||||
stored as that lossless master. MP3 (320 kbps CBR, LAME) is a download
|
||||
format only: `/content?format=mp3` transcodes the FLAC once and caches
|
||||
`{id}.mp3` beside it until the job expires; `response_format` sets a
|
||||
job's default download. Results stored as WAV by older neurons are
|
||||
served as WAV. cortex refuses new jobs (503 `audio_storage_full`) while the job
|
||||
directory's filesystem has under `[audio_jobs] min_free_mb` free (10 GiB
|
||||
default); shared object storage for the Pascal-scale fleet is #411.
|
||||
|
||||
**Metering: per audio-second**, at a per-model rate stored in
|
||||
helexa-upstream's `model_rates` table (`GET /authz/v1/rates`, cached by
|
||||
cortex's `rates.rs`; lookup is the model id, then `'*'`). Unpriced
|
||||
metered submissions get 503. A job settles at the rate it was reserved at.
|
||||
|
||||
**Live numbers (benjy 4090, deployed build, 2026-10-17):**
|
||||
|
||||
| | Resident | 10 s warm | 30 s | Cold swap from text |
|
||||
|---|---|---|---|---|
|
||||
| ACE-Step (planner loaded, off) | ~10.2 GB | 0.56 s | — | ~1 s load |
|
||||
| HeartMuLa | 14.6 GB | 4.3 s (LM 2.4 + codec 1.9) | 11.1 s | 8 s wall incl. load |
|
||||
|
||||
Text restores in 2.5 s. Same seed gives byte-identical output on both
|
||||
engines, and the chain render is byte-identical to the direct neuron
|
||||
render.
|
||||
|
||||
**Sharp edges:**
|
||||
|
||||
- **ACE-Step needs DCW.** The Haar "double" correction after every turbo
|
||||
step: without it, latents are 90% off upstream. Parity with the
|
||||
upstream CPU f32 reference is ~1e-5 with it.
|
||||
- **Seeding is host-side** (SplitMix64 + Box-Muller): candle's CPU
|
||||
`set_seed` is unsupported, and a host RNG makes CPU and GPU agree.
|
||||
- **The 1.7B planner checkpoint has no `model.` prefix; the 0.6B one
|
||||
does.** The loader strips it only where it is absent. A planner that
|
||||
fails to load is non-fatal, so a wrong prefix shows up only as a
|
||||
warning and `plan_ms: 0`.
|
||||
- **The planner matches upstream exactly (#384)** — in f32 with greedy
|
||||
decoding, the CoT byte for byte and every audio code of the trace — and
|
||||
`auto` (the default) now runs it. That parity rests on three things a
|
||||
port gets wrong easily: the CoT grammar (`fsm.rs`, a port of
|
||||
`MetadataConstrainedLogitsProcessor`); the codes prompt, which carries
|
||||
the plan *re-dumped by PyYAML* (`cot.rs`: folded at column 80, quoting
|
||||
values that would read back as bools or numbers); and CFG 2.0 in the
|
||||
codes phase only, against `NO USER INPUT` with empty reasoning.
|
||||
Upstream's own bf16 already picks a different first code, so judge
|
||||
parity in f32.
|
||||
- **ACE-Step's FSQ is the symmetric variant.** `vector_quantize_pytorch`'s
|
||||
`ResidualFSQ` builds its layers with `preserve_symmetry=True` (level k
|
||||
of L → `2k/(L-1) − 1`). The vendored port had the asymmetric formula
|
||||
and tests re-derived from that same formula, so they passed while every
|
||||
planner hint was 18% off. Pin vendored maths to values recorded from
|
||||
the reference library, never to the port's own arithmetic. With codes,
|
||||
upstream also switches the DiT instruction to the cover one.
|
||||
- **Levels follow each upstream.** ACE-Step peak-normalises to −1 dBFS;
|
||||
HeartMuLa does not.
|
||||
- **ACE timing is split** into `encode_ms`, `generate_ms` (denoise) and
|
||||
`decode_ms` (VAE). Each phase ends with a device synchronize, so queued
|
||||
GPU work is charged to the phase that queued it.
|
||||
- **candle only uses its fast transposed-conv path when unpadded.** The
|
||||
Oobleck decoder runs its upsampling unpadded and trims the padding
|
||||
from the output (identical result). Before that, the VAE decode was
|
||||
~95% of a render: 22.8 s for a 30 s song on a 3060, 27 s for 480 s on
|
||||
benjy. Watch for the same trap in any padded `ConvTranspose1d`.
|
||||
- **HeartMuLa's first decode after load is slow** (~33 s, one-time
|
||||
warm-up); every decode after that is ~2 s per 10 s of audio.
|
||||
- **benjy's root drive is slow under pressure** (DRAM-less, 90% full;
|
||||
~8 MB/s flush). First fetches of 20 GB take ~25 min because each
|
||||
shard's close() waits on writeback. Before #409 a client timeout
|
||||
cancelled the load and threw the progress away.
|
||||
- **Validation:** `script/validate-audio.sh [host] [cortex] [model]`
|
||||
(guarded window: evicts and restores co-resident models). helexa-bench
|
||||
has an `audio:<secs>` scenario.
|
||||
|
||||
**Retired:** the PyTorch evaluation beta (`beta.music.helexa.ai`, #388).
|
||||
Its ratings data is archived privately.
|
||||
|
||||
**Layouts (#385).** `[harness.candle.audio]` in neuron.toml sets, per
|
||||
host: `quant = "q8_0"` (DiT, text encoder and planner linears, f32
|
||||
compute), `compute = "f32"`, `placement.<component> = "gpu" | "cpu"`,
|
||||
`planner = "<subfolder>" | "none"`, `max_duration_secs`, and
|
||||
`vram_budget_mb` (smaller VAE chunks). `enforce_vram_budget` holds the
|
||||
rest of the card's memory so a 24 GB card rehearses an 8 GB one.
|
||||
Measured peaks (process, CUDA context included):
|
||||
|
||||
| Layout | Resident | 180 s song | 240 s | 480 s |
|
||||
|---|---|---|---|---|
|
||||
| bf16, no planner | 6.4 GB | 7.8 GB | 7.8 GB | 9.1 GB |
|
||||
| q8 + 8 GiB budget, no planner | 4.6 GB | 6.0 GB | 6.2 GB | 10.0 GB |
|
||||
| q8 + 8 GiB budget + 0.6B planner | 5.8 GB | 7.6 GB | 7.9 GB | 13.1 GB |
|
||||
| dense f32 | 13.2 GB | — | — | — |
|
||||
|
||||
So an 8 GB card serves q8 with the 0.6B planner up to 240 s; beyond
|
||||
that a term grows with the square of song length. q8's mean latent drift
|
||||
from upstream's f32 reference (0.26 rel L2 over four cases) is smaller
|
||||
than upstream's own bf16 (0.32). A 180 s vocal song with the 0.6B
|
||||
planner takes 44 s on a 3060 (q8, 7.1 GB peak). The 0.6B planner is a
|
||||
separate repo, so name it by id: `planner = "ACE-Step/acestep-5Hz-lm-0.6B"`.
|
||||
|
||||
**Next:** the Pascal epic #380 (rig work waits until rigs are named; the fallback if
|
||||
ACE-Step cannot run on Pascal is Stable Audio and/or EzAudio).
|
||||
|
||||
Generated
+91
@@ -311,6 +311,15 @@ version = "1.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
|
||||
|
||||
[[package]]
|
||||
name = "autotools"
|
||||
version = "0.2.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ef941527c41b0fc0dd48511a8154cd5fc7e29200a0ff8b7203c5d777dbc795cf"
|
||||
dependencies = [
|
||||
"cc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-lc-rs"
|
||||
version = "1.16.2"
|
||||
@@ -485,6 +494,12 @@ dependencies = [
|
||||
"tinyvec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "built"
|
||||
version = "0.7.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "56ed6191a7e78c36abdb16ab65341eefd73d64d303fffccdbb00d51e4205967b"
|
||||
|
||||
[[package]]
|
||||
name = "bumpalo"
|
||||
version = "3.20.2"
|
||||
@@ -706,6 +721,12 @@ version = "1.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9"
|
||||
|
||||
[[package]]
|
||||
name = "claxon"
|
||||
version = "0.4.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4bfbf56724aa9eca8afa4fcfadeb479e722935bb2a0900c2d37e0cc477af0688"
|
||||
|
||||
[[package]]
|
||||
name = "cmake"
|
||||
version = "0.1.58"
|
||||
@@ -844,13 +865,17 @@ dependencies = [
|
||||
"axum",
|
||||
"bytes",
|
||||
"chrono",
|
||||
"claxon",
|
||||
"cortex-core",
|
||||
"eventsource-stream",
|
||||
"flacenc",
|
||||
"futures",
|
||||
"helexa-stream",
|
||||
"metrics",
|
||||
"metrics-exporter-prometheus",
|
||||
"mp3lame-encoder",
|
||||
"reqwest",
|
||||
"rustix",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tempfile",
|
||||
@@ -898,6 +923,15 @@ dependencies = [
|
||||
"cfg-if",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crossbeam-channel"
|
||||
version = "0.5.17"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "98b0cc327b5bc766e7fda9c9260cc0fa81b43a8e240440422dff70788e3f9ef1"
|
||||
dependencies = [
|
||||
"crossbeam-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crossbeam-deque"
|
||||
version = "0.8.6"
|
||||
@@ -1479,6 +1513,22 @@ version = "0.5.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99"
|
||||
|
||||
[[package]]
|
||||
name = "flacenc"
|
||||
version = "0.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "74c892c2b5fa08f967e8b5ad29121570b4762202f2402033ce08479ec65eccd0"
|
||||
dependencies = [
|
||||
"built",
|
||||
"crc",
|
||||
"crossbeam-channel",
|
||||
"heapless",
|
||||
"md-5",
|
||||
"num-traits",
|
||||
"rustversion",
|
||||
"seq-macro",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "flagset"
|
||||
version = "0.4.7"
|
||||
@@ -2045,6 +2095,15 @@ dependencies = [
|
||||
"zerocopy",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hash32"
|
||||
version = "0.3.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "47d60b12902ba28e2730cd37e95b8c9223af2808df9e902d4df49588d1470606"
|
||||
dependencies = [
|
||||
"byteorder",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hashbrown"
|
||||
version = "0.12.3"
|
||||
@@ -2108,6 +2167,16 @@ dependencies = [
|
||||
"hashbrown 0.15.5",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "heapless"
|
||||
version = "0.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0bfb9eb618601c89945a70e254898da93b13be0388091d42117462b265bb3fad"
|
||||
dependencies = [
|
||||
"hash32",
|
||||
"stable_deref_trait",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "heck"
|
||||
version = "0.5.0"
|
||||
@@ -3111,6 +3180,26 @@ dependencies = [
|
||||
"pxfm",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "mp3lame-encoder"
|
||||
version = "0.2.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "60cb9bdd89806317373e36ff745f264b7ed7ffc5bc5aab02dc7d1b837c16a8d4"
|
||||
dependencies = [
|
||||
"mp3lame-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "mp3lame-sys"
|
||||
version = "0.1.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "54e3b1772db47828840702e5a2e05694527f731abadf9b931355d54035f019d8"
|
||||
dependencies = [
|
||||
"autotools",
|
||||
"cc",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "native-tls"
|
||||
version = "0.2.18"
|
||||
@@ -3141,10 +3230,12 @@ dependencies = [
|
||||
"candle-nn",
|
||||
"candle-transformers",
|
||||
"clap",
|
||||
"claxon",
|
||||
"cortex-core",
|
||||
"cudaforge",
|
||||
"cudarc 0.19.8",
|
||||
"figment",
|
||||
"flacenc",
|
||||
"futures",
|
||||
"half",
|
||||
"hf-hub",
|
||||
|
||||
@@ -129,3 +129,6 @@ enabled = false
|
||||
# max_attempts = 2
|
||||
# How long a job may wait on a busy model before failing, in seconds.
|
||||
# queue_timeout_secs = 3600
|
||||
# New jobs are refused (503) while the job directory's filesystem has less
|
||||
# than this much free space, in MiB. 0 disables the check.
|
||||
# min_free_mb = 10240
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
//! Keep the cortex binary installable on the oldest glibc we package for.
|
||||
//!
|
||||
//! glibc 2.43 gave `log10f` a new symbol version, so a binary linked on a
|
||||
//! newer host references `log10f@GLIBC_2.43` and will not install on a
|
||||
//! system with an older glibc. The only caller is LAME (the MP3 encoder
|
||||
//! behind `/content?format=mp3`). Routing its calls through
|
||||
//! `__wrap_log10f` (see `glibc_compat.rs`) removes the reference.
|
||||
fn main() {
|
||||
println!("cargo:rustc-link-arg-bins=-Wl,--wrap=log10f");
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
//! Replacements for libm symbols whose current version is newer than the
|
||||
//! glibc cortex is packaged for. `build.rs` points the linker at these.
|
||||
|
||||
/// `log10f` computed through `log10`, which keeps its original symbol
|
||||
/// version. Rounding the f64 result to f32 is at least as accurate as
|
||||
/// libm's single-precision version.
|
||||
#[unsafe(no_mangle)]
|
||||
extern "C" fn __wrap_log10f(x: f32) -> f32 {
|
||||
(x as f64).log10() as f32
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::__wrap_log10f;
|
||||
|
||||
#[test]
|
||||
fn matches_log10_including_edge_cases() {
|
||||
for x in [1e-30_f32, 0.5, 1.0, 2.0, 10.0, 1234.5, f32::MAX] {
|
||||
assert!(
|
||||
(__wrap_log10f(x) - x.log10()).abs() <= 1e-6 * x.log10().abs().max(1.0),
|
||||
"{x}"
|
||||
);
|
||||
}
|
||||
assert_eq!(__wrap_log10f(0.0), f32::NEG_INFINITY);
|
||||
assert!(__wrap_log10f(-1.0).is_nan());
|
||||
assert!(__wrap_log10f(f32::NAN).is_nan());
|
||||
assert_eq!(__wrap_log10f(f32::INFINITY), f32::INFINITY);
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,5 @@
|
||||
mod glibc_compat;
|
||||
|
||||
use anyhow::Result;
|
||||
use clap::{Parser, Subcommand};
|
||||
use cortex_core::config::GatewayConfig;
|
||||
|
||||
@@ -12,10 +12,14 @@
|
||||
//!
|
||||
//! cortex owns the job record. A neuron renders exactly one generation per
|
||||
//! dispatch, synchronously, on `POST /v1/audio/generations`: the response
|
||||
//! body is the audio (`audio/wav`) and the result metadata travels in the
|
||||
//! body is the audio (`audio/flac`) and the result metadata travels in the
|
||||
//! [`RESULT_HEADER`] header as JSON ([`NeuronAudioResult`]). Keeping the
|
||||
//! bytes out of a JSON envelope matters here — ten minutes of 48 kHz stereo
|
||||
//! is over 100 MB of PCM.
|
||||
//!
|
||||
//! Audio is rendered straight to FLAC (16-bit, lossless) and stored that
|
||||
//! way. MP3 is a download format only: cortex transcodes the stored FLAC
|
||||
//! when a client asks for it.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
@@ -44,7 +48,9 @@ pub struct AudioGenerationRequest {
|
||||
/// always reported back, so a random render can be reproduced too.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub seed: Option<u64>,
|
||||
/// Container for the delivered audio. Only `"wav"` today.
|
||||
/// Default download format for the job's content: `"flac"` (the
|
||||
/// default, lossless) or `"mp3"`. The render itself is always FLAC; a
|
||||
/// content fetch can still ask for either with `?format=`.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub response_format: Option<String>,
|
||||
|
||||
@@ -72,7 +78,21 @@ pub struct AudioGenerationRequest {
|
||||
}
|
||||
|
||||
/// Response formats the surface delivers.
|
||||
pub const SUPPORTED_FORMATS: &[&str] = &["wav"];
|
||||
pub const SUPPORTED_FORMATS: &[&str] = &["flac", "mp3"];
|
||||
|
||||
/// The format audio is rendered and stored in.
|
||||
pub const STORED_FORMAT: &str = "flac";
|
||||
|
||||
/// MIME type for an audio format name, including `"wav"` for results stored
|
||||
/// before audio moved to FLAC.
|
||||
pub fn content_type(format: &str) -> &'static str {
|
||||
match format {
|
||||
"flac" => "audio/flac",
|
||||
"mp3" => "audio/mpeg",
|
||||
"wav" => "audio/wav",
|
||||
_ => "application/octet-stream",
|
||||
}
|
||||
}
|
||||
|
||||
impl AudioGenerationRequest {
|
||||
/// True when the request asks for an instrumental (no singing).
|
||||
@@ -150,6 +170,15 @@ pub struct AudioLimits {
|
||||
pub max_steps: Option<u32>,
|
||||
}
|
||||
|
||||
/// A notice an upstream licence asks to be passed on, shown beside the one
|
||||
/// model it applies to (never as a platform-wide licence).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct ModelNotice {
|
||||
pub text: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub url: Option<String>,
|
||||
}
|
||||
|
||||
/// Lifecycle of a job.
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
@@ -213,7 +242,7 @@ pub struct AudioResult {
|
||||
pub channels: u16,
|
||||
/// The seed the render used (random when the request gave none).
|
||||
pub seed: u64,
|
||||
/// Container of the stored audio (`"wav"`).
|
||||
/// Container of the stored audio (`"flac"`).
|
||||
pub format: String,
|
||||
/// Size of the stored audio in bytes.
|
||||
pub bytes: u64,
|
||||
@@ -309,6 +338,10 @@ mod tests {
|
||||
let mut r = req();
|
||||
r.response_format = Some("ogg".into());
|
||||
assert!(r.validate_shape().is_err());
|
||||
r.response_format = Some("wav".into());
|
||||
assert!(r.validate_shape().is_err());
|
||||
r.response_format = Some("mp3".into());
|
||||
assert!(r.validate_shape().is_ok());
|
||||
let mut r = req();
|
||||
r.planner = Some("maybe".into());
|
||||
assert!(r.validate_shape().is_err());
|
||||
|
||||
@@ -94,6 +94,10 @@ pub struct ModelProfile {
|
||||
/// is accepted. `[models.audio]` in TOML.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub audio: Option<crate::audio::AudioLimits>,
|
||||
/// Notices an upstream licence asks to be passed on for this model,
|
||||
/// advertised in `/v1/models` beside it. `[[models.notices]]` in TOML.
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub notices: Vec<crate::audio::ModelNotice>,
|
||||
}
|
||||
|
||||
fn default_min_devices() -> u32 {
|
||||
@@ -289,6 +293,7 @@ mod tests {
|
||||
cost: None,
|
||||
capabilities: vec![],
|
||||
audio: None,
|
||||
notices: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
@@ -537,3 +542,23 @@ source = "helexa"
|
||||
assert_eq!(cat.resolve_alias("helexa/large"), "Qwen/Qwen3.6-27B");
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod repo_catalogue {
|
||||
/// The shipped catalogue parses, and the music model carries its limits.
|
||||
#[test]
|
||||
fn repo_models_toml_parses_with_audio_limits() {
|
||||
let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../../models.toml");
|
||||
let cat = super::ModelCatalogue::load(&path);
|
||||
let ace = cat
|
||||
.get("ACE-Step/Ace-Step1.5")
|
||||
.expect("ACE-Step in models.toml");
|
||||
let audio = ace.audio.as_ref().expect("audio limits");
|
||||
assert!(audio.vocals && audio.max_duration >= audio.default_duration);
|
||||
let heart = cat
|
||||
.get("HeartMuLa/HeartMuLa-oss-3B-happy-new-year")
|
||||
.expect("HeartMuLa in models.toml");
|
||||
assert!(heart.audio.is_some());
|
||||
assert_eq!(heart.notices.len(), 2, "HeartMuLa's notices (#389)");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -57,6 +57,12 @@ pub struct AudioJobsConfig {
|
||||
/// How long a job may wait for a busy model before it fails, in seconds.
|
||||
#[serde(default = "default_audio_queue_timeout")]
|
||||
pub queue_timeout_secs: u64,
|
||||
/// New jobs are refused (503 + `Retry-After`) while the filesystem
|
||||
/// holding `dir` has less than this much free space, in MiB. Rendered
|
||||
/// audio is written there, so this keeps a busy store from filling the
|
||||
/// disk it shares. 0 disables the check.
|
||||
#[serde(default = "default_audio_min_free_mb")]
|
||||
pub min_free_mb: u64,
|
||||
}
|
||||
|
||||
impl Default for AudioJobsConfig {
|
||||
@@ -67,6 +73,7 @@ impl Default for AudioJobsConfig {
|
||||
max_concurrent: default_audio_max_concurrent(),
|
||||
max_attempts: default_audio_max_attempts(),
|
||||
queue_timeout_secs: default_audio_queue_timeout(),
|
||||
min_free_mb: default_audio_min_free_mb(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -86,6 +93,11 @@ fn default_audio_max_attempts() -> u32 {
|
||||
fn default_audio_queue_timeout() -> u64 {
|
||||
3600
|
||||
}
|
||||
/// 10 GiB: room for well over a hundred of the longest songs while jobs
|
||||
/// already accepted finish.
|
||||
fn default_audio_min_free_mb() -> u64 {
|
||||
10 * 1024
|
||||
}
|
||||
|
||||
/// `[upstream]` — the helexa-upstream authority client (#57). Locally
|
||||
/// unrecognised bearer keys are resolved against `url`'s `/authz/v1` surface
|
||||
|
||||
@@ -139,6 +139,11 @@ pub enum ModelStatus {
|
||||
/// retry error instead of 404, and must not race a second
|
||||
/// placement elsewhere.
|
||||
Recovering,
|
||||
/// Gateway-side synthetic status for a model cached on a neuron that
|
||||
/// is currently unhealthy (#376). The location stays visible for
|
||||
/// diagnostics, but nothing will be routed there, so it never makes
|
||||
/// the model count as `loaded`.
|
||||
Unreachable,
|
||||
}
|
||||
|
||||
/// Unified model entry as exposed by the gateway's `/v1/models` endpoint.
|
||||
@@ -194,6 +199,9 @@ pub struct CortexModelEntry {
|
||||
/// audio models with a configured rate.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub audio_tokens_per_second: Option<u64>,
|
||||
/// Notices an upstream licence asks to be passed on for this model.
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub notices: Vec<crate::audio::ModelNotice>,
|
||||
/// `true` when any neuron reports this model supports tool calls.
|
||||
#[serde(default)]
|
||||
pub tool_call: bool,
|
||||
|
||||
@@ -28,7 +28,11 @@ bytes = "1"
|
||||
urlencoding = "2"
|
||||
url = "2"
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
rustix = { version = "1", default-features = false, features = ["std", "fs"] }
|
||||
claxon = "0.4.3"
|
||||
mp3lame-encoder = { version = "0.2.5", features = ["std"] }
|
||||
|
||||
[dev-dependencies]
|
||||
tokio = { workspace = true, features = ["test-util"] }
|
||||
tempfile = "3"
|
||||
flacenc = { version = "0.5.1", default-features = false }
|
||||
|
||||
@@ -0,0 +1,119 @@
|
||||
//! Transcoding stored audio for download.
|
||||
//!
|
||||
//! Renders are stored as 16-bit FLAC, the lossless master. MP3 is offered
|
||||
//! as a download format because some players and users expect it; it is
|
||||
//! encoded from the master on request and cached beside it.
|
||||
|
||||
use anyhow::{Context, bail};
|
||||
use mp3lame_encoder::{Bitrate, Builder, FlushNoGap, InterleavedPcm, Quality};
|
||||
|
||||
/// Frames handed to the MP3 encoder per call, which bounds its output
|
||||
/// buffer for long tracks.
|
||||
const CHUNK_FRAMES: usize = 1152 * 64;
|
||||
|
||||
/// Transcode 16-bit FLAC (mono or stereo) to MP3 at 320 kbps CBR, the
|
||||
/// highest rate MP3 has: it is a copy of a lossless master, so it should
|
||||
/// lose no more than the format forces it to.
|
||||
pub fn flac_to_mp3(flac: &[u8]) -> anyhow::Result<Vec<u8>> {
|
||||
let mut reader =
|
||||
claxon::FlacReader::new(std::io::Cursor::new(flac)).context("reading FLAC header")?;
|
||||
let info = reader.streaminfo();
|
||||
if info.bits_per_sample != 16 {
|
||||
bail!("{}-bit FLAC is not supported", info.bits_per_sample);
|
||||
}
|
||||
let channels = info.channels as usize;
|
||||
if !(1..=2).contains(&channels) {
|
||||
bail!("{channels}-channel FLAC is not supported");
|
||||
}
|
||||
let mut pcm = Vec::with_capacity(info.samples.unwrap_or(0) as usize * channels);
|
||||
for s in reader.samples() {
|
||||
pcm.push(s.context("decoding FLAC")? as i16);
|
||||
}
|
||||
// LAME's interleaved input is stereo; duplicate a mono channel.
|
||||
if channels == 1 {
|
||||
pcm = pcm.iter().flat_map(|&s| [s, s]).collect();
|
||||
}
|
||||
|
||||
let mut builder = Builder::new().context("creating MP3 encoder")?;
|
||||
builder
|
||||
.set_num_channels(2)
|
||||
.map_err(|e| anyhow::anyhow!("MP3 channels: {e}"))?;
|
||||
builder
|
||||
.set_sample_rate(info.sample_rate)
|
||||
.map_err(|e| anyhow::anyhow!("MP3 sample rate: {e}"))?;
|
||||
builder
|
||||
.set_brate(Bitrate::Kbps320)
|
||||
.map_err(|e| anyhow::anyhow!("MP3 bitrate: {e}"))?;
|
||||
builder
|
||||
.set_quality(Quality::NearBest)
|
||||
.map_err(|e| anyhow::anyhow!("MP3 quality: {e}"))?;
|
||||
let mut encoder = builder
|
||||
.build()
|
||||
.map_err(|e| anyhow::anyhow!("MP3 encoder: {e}"))?;
|
||||
|
||||
let mut out = Vec::new();
|
||||
for chunk in pcm.chunks(CHUNK_FRAMES * 2) {
|
||||
out.reserve(mp3lame_encoder::max_required_buffer_size(chunk.len() / 2));
|
||||
encoder
|
||||
.encode_to_vec(InterleavedPcm(chunk), &mut out)
|
||||
.map_err(|e| anyhow::anyhow!("MP3 encode: {e}"))?;
|
||||
}
|
||||
out.reserve(7200);
|
||||
encoder
|
||||
.flush_to_vec::<FlushNoGap>(&mut out)
|
||||
.map_err(|e| anyhow::anyhow!("MP3 flush: {e}"))?;
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn flac(channels: usize, frames: usize) -> Vec<u8> {
|
||||
use flacenc::component::BitRepr;
|
||||
use flacenc::error::Verify;
|
||||
let samples: Vec<i32> = (0..frames * channels)
|
||||
.map(|i| ((i as f32 * 0.03).sin() * 12_000.0) as i32)
|
||||
.collect();
|
||||
let config = flacenc::config::Encoder::default().into_verified().unwrap();
|
||||
let source = flacenc::source::MemSource::from_samples(&samples, channels, 16, 48_000);
|
||||
let stream =
|
||||
flacenc::encode_with_fixed_block_size(&config, source, config.block_size).unwrap();
|
||||
let mut sink = flacenc::bitsink::ByteSink::new();
|
||||
stream.write(&mut sink).unwrap();
|
||||
sink.into_inner()
|
||||
}
|
||||
|
||||
/// The first MPEG audio frame header: sync word, then MPEG-1 Layer III
|
||||
/// at 320 kbps / 48 kHz.
|
||||
fn first_frame(mp3: &[u8]) -> [u8; 4] {
|
||||
let at = mp3
|
||||
.windows(2)
|
||||
.position(|w| w[0] == 0xFF && w[1] & 0xE0 == 0xE0)
|
||||
.expect("no MPEG frame sync");
|
||||
mp3[at..at + 4].try_into().unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stereo_flac_becomes_320k_mp3() {
|
||||
let mp3 = flac_to_mp3(&flac(2, 48_000 * 2)).unwrap();
|
||||
let h = first_frame(&mp3);
|
||||
assert_eq!(h[1] & 0x1E, 0x1A, "not MPEG-1 Layer III: {h:02x?}");
|
||||
assert_eq!(h[2] >> 4, 0xE, "not 320 kbps: {h:02x?}");
|
||||
assert_eq!((h[2] >> 2) & 0x3, 0x1, "not 48 kHz: {h:02x?}");
|
||||
// Two seconds at 320 kbps is 80 kB, give or take a frame.
|
||||
assert!((70_000..90_000).contains(&mp3.len()), "{} bytes", mp3.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mono_flac_is_accepted() {
|
||||
let mp3 = flac_to_mp3(&flac(1, 48_000)).unwrap();
|
||||
assert!(mp3.len() > 30_000, "{} bytes", mp3.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn garbage_is_an_error_not_a_panic() {
|
||||
assert!(flac_to_mp3(b"RIFF....WAVEfmt ").is_err());
|
||||
assert!(flac_to_mp3(&[]).is_err());
|
||||
}
|
||||
}
|
||||
@@ -17,18 +17,21 @@
|
||||
//! account can see, fetch or cancel it. Anonymous jobs (no principal) are
|
||||
//! reachable by id alone — the id is an unguessable capability — and never
|
||||
//! appear in a listing.
|
||||
//!
|
||||
//! Audio is stored as the FLAC the neuron rendered. A content fetch for MP3
|
||||
//! transcodes it once and keeps the MP3 beside it until the job expires.
|
||||
|
||||
use crate::metering::ReservationGuard;
|
||||
use crate::state::CortexState;
|
||||
use axum::Router;
|
||||
use axum::body::{Body, Bytes};
|
||||
use axum::extract::{Path, State};
|
||||
use axum::extract::{Path, Query, State};
|
||||
use axum::http::{HeaderMap, HeaderValue, StatusCode, header};
|
||||
use axum::response::{IntoResponse, Json, Response};
|
||||
use axum::routing::get;
|
||||
use cortex_core::audio::{
|
||||
AudioGenerationRequest, AudioJob, AudioJobList, AudioResult, JOB_OBJECT, JobStatus,
|
||||
NeuronAudioResult, RESULT_HEADER,
|
||||
NeuronAudioResult, RESULT_HEADER, STORED_FORMAT, SUPPORTED_FORMATS,
|
||||
};
|
||||
use cortex_core::config::AudioJobsConfig;
|
||||
use cortex_core::entitlements::{HEADER_ACCOUNT_ID, HEADER_KEY_ID, Principal};
|
||||
@@ -52,6 +55,10 @@ const DEFAULT_BUSY_WAIT: Duration = Duration::from_secs(5);
|
||||
/// How often expired jobs are swept.
|
||||
const SWEEP_INTERVAL: Duration = Duration::from_secs(600);
|
||||
|
||||
/// `Retry-After` on a submission refused for lack of disk space: about as
|
||||
/// long as it takes the sweep to come round and expire old jobs.
|
||||
const STORAGE_RETRY_AFTER_SECS: u64 = 600;
|
||||
|
||||
/// A job as cortex stores it: the public object plus what is needed to run
|
||||
/// (or re-run) it.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
@@ -123,12 +130,37 @@ impl AudioJobs {
|
||||
.as_ref()
|
||||
}
|
||||
|
||||
/// Free space on the filesystem holding the job directory, in MiB.
|
||||
/// `None` when there is no directory (memory only) or it cannot be
|
||||
/// read, in which case submissions are not refused on its account.
|
||||
fn free_mb(&self) -> Option<u64> {
|
||||
let dir = self.dir()?;
|
||||
match rustix::fs::statvfs(dir) {
|
||||
Ok(st) => Some(st.f_bavail.saturating_mul(st.f_frsize) / (1024 * 1024)),
|
||||
Err(e) => {
|
||||
tracing::warn!(dir = %dir.display(), error = %e, "cannot read free space for audio jobs");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether a new job may be accepted: the store's filesystem must keep
|
||||
/// at least `min_free_mb` free. Returns the free space when it may not.
|
||||
fn storage_full(&self) -> Option<u64> {
|
||||
let floor = self.config.min_free_mb;
|
||||
if floor == 0 {
|
||||
return None;
|
||||
}
|
||||
self.free_mb().filter(|free| *free < floor)
|
||||
}
|
||||
|
||||
fn record_path(&self, id: &str) -> Option<PathBuf> {
|
||||
self.dir().map(|d| d.join(format!("{id}.json")))
|
||||
}
|
||||
|
||||
fn audio_path(&self, id: &str) -> Option<PathBuf> {
|
||||
self.dir().map(|d| d.join(format!("{id}.wav")))
|
||||
/// Where a job's audio in `format` lives (or would be cached).
|
||||
fn audio_path(&self, id: &str, format: &str) -> Option<PathBuf> {
|
||||
self.dir().map(|d| d.join(format!("{id}.{format}")))
|
||||
}
|
||||
|
||||
async fn persist(&self, record: &JobRecord) {
|
||||
@@ -270,8 +302,10 @@ pub async fn sweep_expired(fleet: &CortexState) {
|
||||
if let Ok(mut m) = jobs.memory_audio.lock() {
|
||||
m.remove(&id);
|
||||
}
|
||||
for path in [jobs.record_path(&id), jobs.audio_path(&id)]
|
||||
.into_iter()
|
||||
// `wav` covers results stored before audio moved to FLAC.
|
||||
let audio = ["flac", "mp3", "wav"].map(|f| jobs.audio_path(&id, f));
|
||||
for path in std::iter::once(jobs.record_path(&id))
|
||||
.chain(audio)
|
||||
.flatten()
|
||||
{
|
||||
let _ = tokio::fs::remove_file(path).await;
|
||||
@@ -336,6 +370,25 @@ async fn submit(
|
||||
return invalid_params(msg);
|
||||
}
|
||||
|
||||
// Refused before any budget is reserved: nothing to release.
|
||||
if let Some(free) = fleet.audio_jobs.storage_full() {
|
||||
metrics::counter!("cortex_audio_jobs_rejected_total", "reason" => "storage").increment(1);
|
||||
tracing::warn!(
|
||||
free_mb = free,
|
||||
min_free_mb = fleet.audio_jobs.config.min_free_mb,
|
||||
"audio job refused: job store is low on disk space"
|
||||
);
|
||||
return crate::error::envelope_response(
|
||||
OpenAiError::new(
|
||||
503,
|
||||
"server_error",
|
||||
"audio_storage_full",
|
||||
"audio generation is temporarily unavailable: the job store is full",
|
||||
)
|
||||
.with_retry_after(STORAGE_RETRY_AFTER_SECS),
|
||||
);
|
||||
}
|
||||
|
||||
let owner = crate::metering::principal_from_headers(&headers);
|
||||
let mut request = request;
|
||||
request.model = model.clone();
|
||||
@@ -420,11 +473,20 @@ async fn get_job(
|
||||
}
|
||||
}
|
||||
|
||||
/// `GET /v1/audio/generations/{id}/content` — the rendered audio.
|
||||
/// Query string of a content fetch.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ContentQuery {
|
||||
/// `flac` or `mp3`; defaults to the job's `response_format`, then FLAC.
|
||||
format: Option<String>,
|
||||
}
|
||||
|
||||
/// `GET /v1/audio/generations/{id}/content` — the rendered audio, as the
|
||||
/// stored FLAC or transcoded to MP3.
|
||||
async fn get_content(
|
||||
State(fleet): State<Arc<CortexState>>,
|
||||
headers: HeaderMap,
|
||||
Path(id): Path<String>,
|
||||
Query(query): Query<ContentQuery>,
|
||||
) -> Response {
|
||||
let Some(record) = visible_record(&fleet, &headers, &id).await else {
|
||||
return job_not_found(&id);
|
||||
@@ -440,27 +502,70 @@ async fn get_content(
|
||||
),
|
||||
));
|
||||
}
|
||||
let bytes = match fleet.audio_jobs.audio_path(&id) {
|
||||
Some(path) => tokio::fs::read(&path).await.ok(),
|
||||
None => fleet.audio_jobs.memory_audio(&id),
|
||||
let stored = record
|
||||
.job
|
||||
.result
|
||||
.as_ref()
|
||||
.map(|r| r.format.clone())
|
||||
.unwrap_or_else(|| STORED_FORMAT.into());
|
||||
let wanted = query
|
||||
.format
|
||||
.or_else(|| record.request.response_format.clone())
|
||||
.unwrap_or_else(|| stored.clone());
|
||||
if wanted != stored && !SUPPORTED_FORMATS.contains(&wanted.as_str()) {
|
||||
return invalid_params(format!(
|
||||
"format '{wanted}' is not supported; use one of {SUPPORTED_FORMATS:?}"
|
||||
));
|
||||
}
|
||||
|
||||
let Some(master) = fleet.audio_jobs.load_audio(&id, &stored).await else {
|
||||
return content_expired(&id);
|
||||
};
|
||||
let Some(bytes) = bytes else {
|
||||
let bytes = if wanted == stored {
|
||||
master
|
||||
} else if wanted == "mp3" && stored == "flac" {
|
||||
match fleet.audio_jobs.mp3(&id, master).await {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
tracing::warn!(job = %id, error = %e, "mp3 transcode failed");
|
||||
return crate::error::envelope_response(OpenAiError::new(
|
||||
500,
|
||||
"api_error",
|
||||
"transcode_failed",
|
||||
format!("could not produce mp3 for audio job '{id}'"),
|
||||
));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Only results stored before the move to FLAC get here.
|
||||
return crate::error::envelope_response(OpenAiError::new(
|
||||
410,
|
||||
409,
|
||||
"invalid_request_error",
|
||||
"content_expired",
|
||||
format!("the audio for job '{id}' is no longer stored"),
|
||||
"format_unavailable",
|
||||
format!("audio job '{id}' is stored as {stored} and is only available as {stored}"),
|
||||
));
|
||||
};
|
||||
|
||||
let mut resp = Response::new(Body::from(bytes));
|
||||
resp.headers_mut()
|
||||
.insert(header::CONTENT_TYPE, HeaderValue::from_static("audio/wav"));
|
||||
if let Ok(v) = HeaderValue::from_str(&format!("attachment; filename=\"{id}.wav\"")) {
|
||||
resp.headers_mut().insert(
|
||||
header::CONTENT_TYPE,
|
||||
HeaderValue::from_static(cortex_core::audio::content_type(&wanted)),
|
||||
);
|
||||
if let Ok(v) = HeaderValue::from_str(&format!("attachment; filename=\"{id}.{wanted}\"")) {
|
||||
resp.headers_mut().insert(header::CONTENT_DISPOSITION, v);
|
||||
}
|
||||
resp
|
||||
}
|
||||
|
||||
fn content_expired(id: &str) -> Response {
|
||||
crate::error::envelope_response(OpenAiError::new(
|
||||
410,
|
||||
"invalid_request_error",
|
||||
"content_expired",
|
||||
format!("the audio for job '{id}' is no longer stored"),
|
||||
))
|
||||
}
|
||||
|
||||
/// `DELETE /v1/audio/generations/{id}` — cancel a queued or running job.
|
||||
/// Cancelling a finished job is a no-op that returns it unchanged.
|
||||
async fn cancel_job(
|
||||
@@ -530,7 +635,8 @@ fn spawn_worker(fleet: Arc<CortexState>, id: String, guard: ReservationGuard) {
|
||||
|
||||
/// Outcome of one dispatch to a neuron.
|
||||
enum Dispatch {
|
||||
Done(Bytes, NeuronAudioResult),
|
||||
/// The rendered audio, its format, and the result metadata.
|
||||
Done(Bytes, String, NeuronAudioResult),
|
||||
/// The model is busy; try again after the given wait.
|
||||
Busy(Duration),
|
||||
/// The request itself was refused; retrying will not help.
|
||||
@@ -571,8 +677,8 @@ async fn run_job(fleet: &Arc<CortexState>, id: &str, guard: ReservationGuard) {
|
||||
while attempts < jobs.config.max_attempts.max(1) {
|
||||
let outcome = dispatch_once(fleet, &record).await;
|
||||
match outcome {
|
||||
Dispatch::Done(bytes, result) => {
|
||||
complete(fleet, id, &record, bytes, result, guard).await;
|
||||
Dispatch::Done(bytes, format, result) => {
|
||||
complete(fleet, id, &record, bytes, &format, result, guard).await;
|
||||
return;
|
||||
}
|
||||
Dispatch::Busy(wait) => {
|
||||
@@ -659,6 +765,15 @@ async fn dispatch_once(fleet: &Arc<CortexState>, record: &JobRecord) -> Dispatch
|
||||
.get(RESULT_HEADER)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(str::to_string);
|
||||
// Neurons render FLAC; one still on an older build sends WAV.
|
||||
let format = match resp
|
||||
.headers()
|
||||
.get(header::CONTENT_TYPE)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
{
|
||||
Some(t) if t.starts_with("audio/wav") => "wav",
|
||||
_ => STORED_FORMAT,
|
||||
};
|
||||
let bytes = match resp.bytes().await {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
@@ -678,7 +793,7 @@ async fn dispatch_once(fleet: &Arc<CortexState>, record: &JobRecord) -> Dispatch
|
||||
.as_deref()
|
||||
.and_then(|h| serde_json::from_str::<NeuronAudioResult>(h).ok());
|
||||
return match parsed {
|
||||
Some(result) => Dispatch::Done(bytes, result),
|
||||
Some(result) => Dispatch::Done(bytes, format.to_string(), result),
|
||||
None => Dispatch::Failed(error_value(
|
||||
"api_error",
|
||||
"upstream_bad_response",
|
||||
@@ -712,6 +827,7 @@ async fn complete(
|
||||
id: &str,
|
||||
record: &JobRecord,
|
||||
bytes: Bytes,
|
||||
format: &str,
|
||||
result: NeuronAudioResult,
|
||||
guard: ReservationGuard,
|
||||
) {
|
||||
@@ -723,7 +839,7 @@ async fn complete(
|
||||
return;
|
||||
}
|
||||
let size = bytes.len() as u64;
|
||||
match jobs.audio_path(id) {
|
||||
match jobs.audio_path(id, format) {
|
||||
Some(path) => {
|
||||
if let Err(e) = tokio::fs::write(&path, &bytes).await {
|
||||
drop(guard);
|
||||
@@ -762,7 +878,7 @@ async fn complete(
|
||||
sample_rate: result.sample_rate,
|
||||
channels: result.channels,
|
||||
seed: result.seed,
|
||||
format: "wav".into(),
|
||||
format: format.to_string(),
|
||||
bytes: size,
|
||||
});
|
||||
r.job.usage = Some(result.usage.clone());
|
||||
@@ -837,6 +953,37 @@ impl AudioJobs {
|
||||
.ok()
|
||||
.and_then(|m| m.get(id).map(|b| b.to_vec()))
|
||||
}
|
||||
|
||||
/// A job's stored audio, from disk or memory.
|
||||
async fn load_audio(&self, id: &str, format: &str) -> Option<Vec<u8>> {
|
||||
match self.audio_path(id, format) {
|
||||
Some(path) => tokio::fs::read(&path).await.ok(),
|
||||
None => self.memory_audio(id),
|
||||
}
|
||||
}
|
||||
|
||||
/// The MP3 of a job's FLAC: the cached copy if there is one, else a
|
||||
/// fresh transcode, cached when there is a job directory.
|
||||
async fn mp3(&self, id: &str, flac: Vec<u8>) -> anyhow::Result<Vec<u8>> {
|
||||
let cache = self.audio_path(id, "mp3");
|
||||
if let Some(path) = &cache
|
||||
&& let Ok(bytes) = tokio::fs::read(path).await
|
||||
{
|
||||
return Ok(bytes);
|
||||
}
|
||||
let mp3 =
|
||||
tokio::task::spawn_blocking(move || crate::audio_codec::flac_to_mp3(&flac)).await??;
|
||||
if let Some(path) = cache {
|
||||
// Write-then-rename so a concurrent fetch never reads a torn file.
|
||||
let tmp = path.with_extension("mp3.tmp");
|
||||
let written = tokio::fs::write(&tmp, &mp3).await.is_ok()
|
||||
&& tokio::fs::rename(&tmp, &path).await.is_ok();
|
||||
if !written {
|
||||
tracing::warn!(job = %id, "could not cache the mp3 transcode");
|
||||
}
|
||||
}
|
||||
Ok(mp3)
|
||||
}
|
||||
}
|
||||
|
||||
fn now_secs() -> u64 {
|
||||
|
||||
@@ -46,6 +46,7 @@ pub async fn evict_lru_on_node(
|
||||
.models
|
||||
.values()
|
||||
.filter(|m| m.status == ModelStatus::Loaded)
|
||||
.filter(|m| !fleet.eviction_blocked(node, &m.id))
|
||||
.filter(|m| match incoming {
|
||||
Some(inc) => fleet.catalogue.may_displace(inc, &m.id),
|
||||
None => {
|
||||
|
||||
@@ -980,6 +980,7 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
cost: profile.cost.clone(),
|
||||
audio: profile.audio.clone(),
|
||||
audio_tokens_per_second: None,
|
||||
notices: profile.notices.clone(),
|
||||
// Runtime-detected — will be OR-ed in Pass 2 from neuron data.
|
||||
tool_call: false,
|
||||
reasoning: false,
|
||||
@@ -1001,14 +1002,24 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
// in the catalogue, create a new CortexModelEntry from scratch —
|
||||
// cortex doesn't refuse to surface a manually-loaded model just
|
||||
// because the operator didn't enumerate it in models.toml.
|
||||
//
|
||||
// An unhealthy neuron's models are its last snapshot, not live
|
||||
// state, and the router will not send it anything (#376). They are
|
||||
// listed as `unreachable` and never make a model `loaded`, so
|
||||
// `loaded` means the same thing to a client as it does to routing.
|
||||
for node in nodes.values() {
|
||||
for (model_id, entry) in &node.models {
|
||||
let status = if node.healthy {
|
||||
entry.status
|
||||
} else {
|
||||
cortex_core::node::ModelStatus::Unreachable
|
||||
};
|
||||
let location = ModelLocation {
|
||||
node: node.name.clone(),
|
||||
status: entry.status,
|
||||
status,
|
||||
vram_estimate_mb: entry.vram_estimate_mb,
|
||||
};
|
||||
let was_loaded = matches!(entry.status, cortex_core::node::ModelStatus::Loaded);
|
||||
let was_loaded = matches!(status, cortex_core::node::ModelStatus::Loaded);
|
||||
entries
|
||||
.entry(model_id.clone())
|
||||
.and_modify(|e| {
|
||||
@@ -1064,6 +1075,7 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
cost: None,
|
||||
audio: None,
|
||||
audio_tokens_per_second: None,
|
||||
notices: Vec::new(),
|
||||
tool_call: entry.tool_call,
|
||||
reasoning: entry.reasoning,
|
||||
max_model_len: None,
|
||||
@@ -1091,7 +1103,9 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
// Synthesising a Loading location instead tells clients the model
|
||||
// is on its way. Idempotent against Pass 2: if a Loading location
|
||||
// for this node already exists (shouldn't, but be safe) we skip.
|
||||
for node in nodes.values() {
|
||||
// An unhealthy neuron's activation snapshot is as stale as its
|
||||
// models, so it reports nothing as on its way.
|
||||
for node in nodes.values().filter(|n| n.healthy) {
|
||||
let Some(activation) = node.activation.as_ref() else {
|
||||
continue;
|
||||
};
|
||||
@@ -1135,6 +1149,7 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
cost: None,
|
||||
audio: None,
|
||||
audio_tokens_per_second: None,
|
||||
notices: Vec::new(),
|
||||
tool_call: false,
|
||||
reasoning: false,
|
||||
max_model_len: None,
|
||||
@@ -1178,6 +1193,7 @@ async fn list_models(State(fleet): State<Arc<CortexState>>) -> Json<Value> {
|
||||
cost: target_entry.cost.clone(),
|
||||
audio: target_entry.audio.clone(),
|
||||
audio_tokens_per_second: None,
|
||||
notices: target_entry.notices.clone(),
|
||||
tool_call: target_entry.tool_call,
|
||||
reasoning: target_entry.reasoning,
|
||||
max_model_len: None,
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
pub mod anthropic_sse;
|
||||
pub mod audio_codec;
|
||||
pub mod audio_jobs;
|
||||
pub mod auth;
|
||||
pub mod entitlements_chain;
|
||||
|
||||
@@ -310,7 +310,7 @@ async fn discover_capabilities(fleet: &CortexState, name: &str, endpoint: &str)
|
||||
/// Fetch `/health` and stash the activation snapshot on NodeState.
|
||||
/// Decoupled from the /models poll so a /health glitch doesn't mark
|
||||
/// the neuron unhealthy or evict the model list.
|
||||
async fn poll_health(fleet: &CortexState, name: &str, endpoint: &str) {
|
||||
pub async fn poll_health(fleet: &CortexState, name: &str, endpoint: &str) {
|
||||
let url = format!("{endpoint}/health");
|
||||
let resp = match fleet
|
||||
.http_client
|
||||
|
||||
@@ -64,6 +64,10 @@ pub enum RouteError {
|
||||
"model '{model_id}' is recovering on node '{node}' (device context rebuild in progress) — retry shortly"
|
||||
)]
|
||||
ModelRecovering { model_id: String, node: String },
|
||||
#[error(
|
||||
"model '{model_id}' does not fit on '{node}' until a model there finishes its work — retry shortly"
|
||||
)]
|
||||
NodeBusy { model_id: String, node: String },
|
||||
}
|
||||
|
||||
impl RouteError {
|
||||
@@ -74,7 +78,8 @@ impl RouteError {
|
||||
match self {
|
||||
RouteError::NoHealthyNodes
|
||||
| RouteError::ModelRecovering { .. }
|
||||
| RouteError::FeasibleNodeUnhealthy { .. } => 503,
|
||||
| RouteError::FeasibleNodeUnhealthy { .. }
|
||||
| RouteError::NodeBusy { .. } => 503,
|
||||
_ => 404,
|
||||
}
|
||||
}
|
||||
@@ -88,7 +93,8 @@ impl RouteError {
|
||||
| RouteError::NoFeasibleNeuron { .. }
|
||||
| RouteError::ColdLoadFailed { .. }
|
||||
| RouteError::ModelRecovering { .. }
|
||||
| RouteError::FeasibleNodeUnhealthy { .. } => "api_error",
|
||||
| RouteError::FeasibleNodeUnhealthy { .. }
|
||||
| RouteError::NodeBusy { .. } => "api_error",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -102,6 +108,7 @@ impl RouteError {
|
||||
RouteError::ColdLoadFailed { .. } => "service_unavailable",
|
||||
RouteError::ModelRecovering { .. } => "service_unavailable",
|
||||
RouteError::FeasibleNodeUnhealthy { .. } => "service_unavailable",
|
||||
RouteError::NodeBusy { .. } => "service_unavailable",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -115,6 +122,8 @@ impl RouteError {
|
||||
RouteError::ModelRecovering { .. } => Some(2),
|
||||
RouteError::FeasibleNodeUnhealthy { .. } => Some(3),
|
||||
RouteError::NoHealthyNodes => Some(5),
|
||||
// Clears when the resident model's requests finish.
|
||||
RouteError::NodeBusy { .. } => Some(10),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
@@ -205,12 +214,14 @@ pub async fn resolve(
|
||||
// Loading is gateway-synthesised from neuron's
|
||||
// activation snapshot; it never appears on the
|
||||
// wire from neuron's `/models`. Skip — the model
|
||||
// isn't actually servable yet. The pre-existing
|
||||
// race (catalogue cold_load fires a parallel
|
||||
// /models/load against the in-flight load) is no
|
||||
// worse than before; fixing it needs neuron-side
|
||||
// in-flight tracking on /models/load itself.
|
||||
// isn't actually servable yet. Cortex's own
|
||||
// cold-loads are serialised per node (priority 4);
|
||||
// a load neuron started by itself at activation is
|
||||
// still invisible to that lock.
|
||||
ModelStatus::Loading => {}
|
||||
// Synthesised only by `/v1/models` for unhealthy
|
||||
// neurons, which this loop has already skipped.
|
||||
ModelStatus::Unreachable => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -247,41 +258,109 @@ pub async fn resolve(
|
||||
|
||||
// Priority 4: catalogue × topology cold-load.
|
||||
if let Some(profile) = fleet.catalogue.get(model_id) {
|
||||
let (node_name, neuron_endpoint, fits_free) = pick_feasible_neuron(fleet, profile).await?;
|
||||
// Cold-swap (#203): when the chosen node's devices don't have the
|
||||
// free VRAM *right now*, evict unpinned LRU models there until
|
||||
// the profile fits (bounded — a node whose evictable set can't
|
||||
// free enough was ranked below one that can).
|
||||
if !fits_free {
|
||||
for _ in 0..3 {
|
||||
match crate::evictor::evict_lru_on_node(fleet, &node_name, Some(&profile.id)).await
|
||||
{
|
||||
Ok(Some(evicted)) => {
|
||||
tracing::info!(
|
||||
model = %profile.id,
|
||||
node = %node_name,
|
||||
evicted = %evicted,
|
||||
"cold-swap: evicted LRU model to make room"
|
||||
);
|
||||
if node_fits_free(fleet, &node_name, profile).await {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(None) => break,
|
||||
Err(e) => {
|
||||
tracing::warn!(node = %node_name, error = %e, "cold-swap eviction failed");
|
||||
break;
|
||||
}
|
||||
let (node_name, neuron_endpoint, _) = pick_feasible_neuron(fleet, profile).await?;
|
||||
// One cold-load per node at a time. The choice above rests on the
|
||||
// last poll's VRAM, which cannot see a load in progress, so every
|
||||
// decision that follows is re-made inside the lock from a fresh
|
||||
// `/health`.
|
||||
let lock = fleet.cold_load_lock(&node_name);
|
||||
let _turn = lock.lock().await;
|
||||
|
||||
// Whoever held the lock may have been loading this very model.
|
||||
if let Some((node, endpoint)) = loaded_location(fleet, model_id).await {
|
||||
return finish(fleet, &node, &endpoint, model_id, false).await;
|
||||
}
|
||||
|
||||
crate::poller::poll_health(fleet, &node_name, &neuron_endpoint).await;
|
||||
// Cold-swap (#203): when the node's devices don't have the free
|
||||
// VRAM right now, evict unpinned LRU models there until the
|
||||
// profile fits (bounded — a node whose evictable set can't free
|
||||
// enough was ranked below one that can).
|
||||
for _ in 0..3 {
|
||||
if node_fits_free(fleet, &node_name, profile).await {
|
||||
break;
|
||||
}
|
||||
match crate::evictor::evict_lru_on_node(fleet, &node_name, Some(&profile.id)).await {
|
||||
Ok(Some(evicted)) => {
|
||||
tracing::info!(
|
||||
model = %profile.id,
|
||||
node = %node_name,
|
||||
evicted = %evicted,
|
||||
"cold-swap: evicted LRU model to make room"
|
||||
);
|
||||
crate::poller::poll_health(fleet, &node_name, &neuron_endpoint).await;
|
||||
}
|
||||
Ok(None) => break,
|
||||
Err(e) => {
|
||||
tracing::warn!(node = %node_name, error = %e, "cold-swap eviction failed");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
// Still short, and a model that would otherwise be displaced is
|
||||
// busy: loading now would land beside it and fail out of memory.
|
||||
// Answer "retry shortly" instead; the room appears when its work
|
||||
// ends.
|
||||
if !node_fits_free(fleet, &node_name, profile).await
|
||||
&& blocked_displaceable(fleet, &node_name, profile).await
|
||||
{
|
||||
tracing::info!(
|
||||
model = %profile.id,
|
||||
node = %node_name,
|
||||
"cold-load deferred: the model it would displace is busy"
|
||||
);
|
||||
return Err(RouteError::NodeBusy {
|
||||
model_id: profile.id.clone(),
|
||||
node: node_name,
|
||||
});
|
||||
}
|
||||
cold_load(fleet, &node_name, &neuron_endpoint, profile).await?;
|
||||
fleet.note_cold_load(&node_name, &profile.id);
|
||||
// The next cold-load in line decides from these readings.
|
||||
crate::poller::poll_health(fleet, &node_name, &neuron_endpoint).await;
|
||||
drop(_turn);
|
||||
return finish(fleet, &node_name, &neuron_endpoint, model_id, true).await;
|
||||
}
|
||||
|
||||
Err(RouteError::ModelNotFound(model_id.to_string()))
|
||||
}
|
||||
|
||||
/// A healthy node where `model_id` is loaded and servable, if any.
|
||||
async fn loaded_location(fleet: &Arc<CortexState>, model_id: &str) -> Option<(String, String)> {
|
||||
let nodes = fleet.nodes.read().await;
|
||||
nodes
|
||||
.values()
|
||||
.filter(|n| n.healthy)
|
||||
.find(|n| {
|
||||
n.models.get(model_id).is_some_and(|m| {
|
||||
matches!(m.status, ModelStatus::Loaded | ModelStatus::Reloading) && m.is_servable()
|
||||
})
|
||||
})
|
||||
.map(|n| (n.name.clone(), n.endpoint.clone()))
|
||||
}
|
||||
|
||||
/// Whether `node_name` holds a model the profile may displace but that
|
||||
/// cannot be evicted right now because it is busy (see
|
||||
/// [`CortexState::eviction_blocked`]), or is still loading its startup
|
||||
/// models.
|
||||
async fn blocked_displaceable(
|
||||
fleet: &Arc<CortexState>,
|
||||
node_name: &str,
|
||||
profile: &ModelProfile,
|
||||
) -> bool {
|
||||
let nodes = fleet.nodes.read().await;
|
||||
nodes.get(node_name).is_some_and(|node| {
|
||||
// A startup load cannot be evicted until it finishes, and its
|
||||
// memory is still arriving: wait for it.
|
||||
activating(node)
|
||||
|| node.models.values().any(|m| {
|
||||
m.status == ModelStatus::Loaded
|
||||
&& fleet.catalogue.may_displace(&profile.id, &m.id)
|
||||
&& fleet.eviction_blocked(node, &m.id)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
/// True when `node_name`'s live device readings show enough free VRAM
|
||||
/// on some device for the profile. Conservative: unknown health (no
|
||||
/// /health poll yet) counts as not-fitting so the ranking prefers
|
||||
@@ -289,12 +368,40 @@ pub async fn resolve(
|
||||
async fn node_fits_free(fleet: &Arc<CortexState>, node_name: &str, profile: &ModelProfile) -> bool {
|
||||
let nodes = fleet.nodes.read().await;
|
||||
nodes.get(node_name).is_some_and(|node| {
|
||||
let reserved = activation_reserved_mb(fleet, node);
|
||||
node.device_health
|
||||
.iter()
|
||||
.any(|d| d.vram_free_mb >= profile.vram_mb.unwrap_or(0))
|
||||
.any(|d| d.vram_free_mb.saturating_sub(reserved) >= profile.vram_mb.unwrap_or(0))
|
||||
})
|
||||
}
|
||||
|
||||
/// VRAM promised to models the neuron is loading on its own at startup
|
||||
/// (its `default_models`), by their catalogue footprint. Those loads are
|
||||
/// not in the neuron's model list until they finish and their memory
|
||||
/// fills in gradually, so free VRAM alone overstates what is available.
|
||||
fn activation_reserved_mb(fleet: &CortexState, node: &cortex_core::node::NodeState) -> u64 {
|
||||
node.activation.as_ref().map_or(0, |a| {
|
||||
a.in_progress
|
||||
.iter()
|
||||
.chain(a.pending.iter())
|
||||
.filter(|id| {
|
||||
!node
|
||||
.models
|
||||
.get(id.as_str())
|
||||
.is_some_and(|m| m.status == ModelStatus::Loaded)
|
||||
})
|
||||
.filter_map(|id| fleet.catalogue.get(id).and_then(|p| p.vram_mb))
|
||||
.sum()
|
||||
})
|
||||
}
|
||||
|
||||
/// Whether the neuron is still loading its startup models.
|
||||
fn activating(node: &cortex_core::node::NodeState) -> bool {
|
||||
node.activation
|
||||
.as_ref()
|
||||
.is_some_and(|a| a.in_progress.is_some() || !a.pending.is_empty())
|
||||
}
|
||||
|
||||
/// Pick a healthy neuron whose discovered topology satisfies the
|
||||
/// profile. Preference order (#203):
|
||||
/// 1. A neuron from `profile.pinned_on` that is healthy + feasible.
|
||||
@@ -332,7 +439,8 @@ async fn pick_feasible_neuron(
|
||||
.iter()
|
||||
.map(|d| d.vram_free_mb)
|
||||
.max()
|
||||
.unwrap_or(0);
|
||||
.unwrap_or(0)
|
||||
.saturating_sub(activation_reserved_mb(fleet, node));
|
||||
let need = profile.vram_mb.unwrap_or(0);
|
||||
let fits_free = max_free >= need;
|
||||
// Evictable estimate: free + the VRAM of loaded models this
|
||||
@@ -344,6 +452,7 @@ async fn pick_feasible_neuron(
|
||||
.values()
|
||||
.filter(|m| {
|
||||
matches!(m.status, cortex_core::node::ModelStatus::Loaded)
|
||||
&& !fleet.eviction_blocked(node, &m.id)
|
||||
&& fleet.catalogue.may_displace(&profile.id, &m.id)
|
||||
})
|
||||
// neuron reports vram_used_mb: null today; fall back to the
|
||||
@@ -664,6 +773,7 @@ mod tests {
|
||||
cost: None,
|
||||
capabilities: vec![],
|
||||
audio: None,
|
||||
notices: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -57,8 +57,22 @@ pub struct CortexState {
|
||||
/// Per-model rates for non-token units (#386), from upstream when it
|
||||
/// is enabled, else the local default.
|
||||
pub rates: Arc<crate::rates::Rates>,
|
||||
/// One lock per node, held across a cold-load's evict → load → fresh
|
||||
/// `/health` sequence. Placement reads VRAM from the last poll, which
|
||||
/// cannot see a load still in progress; two cold-loads deciding at
|
||||
/// once both saw the same free card and loaded side by side into an
|
||||
/// out-of-memory.
|
||||
cold_load_locks: std::sync::Mutex<HashMap<String, Arc<tokio::sync::Mutex<()>>>>,
|
||||
/// When cortex last cold-loaded each `(node, model)`. A model loaded
|
||||
/// for a request is protected from eviction for [`COLD_LOAD_GRACE`],
|
||||
/// long enough for that request to reach it; otherwise the next
|
||||
/// cold-load in line evicts it before it has served anything.
|
||||
recent_cold_loads: std::sync::Mutex<HashMap<(String, String), std::time::Instant>>,
|
||||
}
|
||||
|
||||
/// How long a freshly cold-loaded model is protected from eviction.
|
||||
pub const COLD_LOAD_GRACE: std::time::Duration = std::time::Duration::from_secs(30);
|
||||
|
||||
impl CortexState {
|
||||
pub fn from_config(config: &GatewayConfig) -> Self {
|
||||
let mut nodes = HashMap::new();
|
||||
@@ -142,6 +156,53 @@ impl CortexState {
|
||||
} else {
|
||||
crate::rates::Rates::local()
|
||||
}),
|
||||
cold_load_locks: std::sync::Mutex::new(HashMap::new()),
|
||||
recent_cold_loads: std::sync::Mutex::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// The lock that serialises cold-loads on `node`.
|
||||
pub fn cold_load_lock(&self, node: &str) -> Arc<tokio::sync::Mutex<()>> {
|
||||
let mut locks = self
|
||||
.cold_load_locks
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
Arc::clone(locks.entry(node.to_string()).or_default())
|
||||
}
|
||||
|
||||
/// Record that cortex has just cold-loaded `model` on `node`.
|
||||
pub fn note_cold_load(&self, node: &str, model: &str) {
|
||||
let mut recent = self
|
||||
.recent_cold_loads
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
recent.retain(|_, at| at.elapsed() < COLD_LOAD_GRACE);
|
||||
recent.insert(
|
||||
(node.to_string(), model.to_string()),
|
||||
std::time::Instant::now(),
|
||||
);
|
||||
}
|
||||
|
||||
/// Whether `model` on `node` is inside its post-cold-load grace.
|
||||
pub fn recently_cold_loaded(&self, node: &str, model: &str) -> bool {
|
||||
let recent = self
|
||||
.recent_cold_loads
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
recent
|
||||
.get(&(node.to_string(), model.to_string()))
|
||||
.is_some_and(|at| at.elapsed() < COLD_LOAD_GRACE)
|
||||
}
|
||||
|
||||
/// Whether a resident model must not be evicted right now: it has
|
||||
/// requests running or queued (as of the last `/health`), or cortex
|
||||
/// cold-loaded it moments ago for a request that has yet to arrive.
|
||||
/// Evicting either kills work in progress.
|
||||
pub fn eviction_blocked(&self, node: &cortex_core::node::NodeState, model: &str) -> bool {
|
||||
let busy = node
|
||||
.model_load
|
||||
.get(model)
|
||||
.is_some_and(|l| l.in_flight + l.queue_depth > 0);
|
||||
busy || self.recently_cold_loaded(&node.name, model)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -22,8 +22,25 @@ use std::time::Duration;
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
const MODEL: &str = "music-model";
|
||||
/// What a neuron on an older build rendered: WAV.
|
||||
const WAV: &[u8] = b"RIFF\x24\x00\x00\x00WAVEfmt fake-audio-bytes";
|
||||
|
||||
/// What a neuron renders: half a second of a real 16-bit stereo FLAC, so
|
||||
/// that cortex can transcode it.
|
||||
static FLAC: std::sync::LazyLock<Vec<u8>> = std::sync::LazyLock::new(|| {
|
||||
use flacenc::component::BitRepr;
|
||||
use flacenc::error::Verify;
|
||||
let samples: Vec<i32> = (0..48_000)
|
||||
.map(|i| ((i as f32 * 0.02).sin() * 8_000.0) as i32)
|
||||
.collect();
|
||||
let config = flacenc::config::Encoder::default().into_verified().unwrap();
|
||||
let source = flacenc::source::MemSource::from_samples(&samples, 2, 16, 48_000);
|
||||
let stream = flacenc::encode_with_fixed_block_size(&config, source, config.block_size).unwrap();
|
||||
let mut sink = flacenc::bitsink::ByteSink::new();
|
||||
stream.write(&mut sink).unwrap();
|
||||
sink.into_inner()
|
||||
});
|
||||
|
||||
/// How the mock neuron answers `POST /v1/audio/generations`.
|
||||
#[derive(Clone, Copy)]
|
||||
enum Mode {
|
||||
@@ -41,6 +58,8 @@ enum Mode {
|
||||
Slow(Duration),
|
||||
/// Always 503.
|
||||
AlwaysBusy,
|
||||
/// Render WAV, as neurons did before audio moved to FLAC.
|
||||
LegacyWav,
|
||||
}
|
||||
|
||||
struct Mock {
|
||||
@@ -136,13 +155,17 @@ async fn render(
|
||||
helexa_timing: None,
|
||||
},
|
||||
};
|
||||
let (content_type, body) = match mode {
|
||||
Mode::LegacyWav => ("audio/wav", WAV.to_vec()),
|
||||
_ => ("audio/flac", FLAC.clone()),
|
||||
};
|
||||
(
|
||||
StatusCode::OK,
|
||||
[
|
||||
("content-type", "audio/wav".to_string()),
|
||||
("content-type", content_type.to_string()),
|
||||
(RESULT_HEADER, serde_json::to_string(&result).unwrap()),
|
||||
],
|
||||
WAV,
|
||||
body,
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
@@ -179,6 +202,9 @@ fn config(neuron: &str, dir: &std::path::Path) -> GatewayConfig {
|
||||
audio_jobs: AudioJobsConfig {
|
||||
dir: dir.to_string_lossy().into_owned(),
|
||||
queue_timeout_secs: 20,
|
||||
// The free-space floor is tested on its own below; elsewhere it
|
||||
// must not make results depend on the build host's disk.
|
||||
min_free_mb: 0,
|
||||
..Default::default()
|
||||
},
|
||||
}
|
||||
@@ -283,7 +309,8 @@ async fn submit_poll_fetch_content() {
|
||||
assert_eq!(done["result"]["duration"], 12.5);
|
||||
assert_eq!(done["result"]["sample_rate"], 48_000);
|
||||
assert_eq!(done["result"]["seed"], 42);
|
||||
assert_eq!(done["result"]["bytes"], WAV.len());
|
||||
assert_eq!(done["result"]["bytes"], FLAC.len());
|
||||
assert_eq!(done["result"]["format"], "flac");
|
||||
assert_eq!(done["usage"]["helexa_audio_units"], 12.5);
|
||||
|
||||
let content = reqwest::Client::new()
|
||||
@@ -293,8 +320,12 @@ async fn submit_poll_fetch_content() {
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(content.status(), 200);
|
||||
assert_eq!(content.headers()["content-type"], "audio/wav");
|
||||
assert_eq!(content.bytes().await.unwrap().as_ref(), WAV);
|
||||
assert_eq!(content.headers()["content-type"], "audio/flac");
|
||||
assert_eq!(
|
||||
content.headers()["content-disposition"],
|
||||
format!("attachment; filename=\"{id}.flac\"").as_str()
|
||||
);
|
||||
assert_eq!(content.bytes().await.unwrap().as_ref(), FLAC.as_slice());
|
||||
|
||||
// The neuron saw the submitting account, for fair-share admission.
|
||||
assert_eq!(
|
||||
@@ -303,7 +334,110 @@ async fn submit_poll_fetch_content() {
|
||||
);
|
||||
// Record and audio are on disk.
|
||||
assert!(dir.path().join(format!("{id}.json")).exists());
|
||||
assert!(dir.path().join(format!("{id}.flac")).exists());
|
||||
}
|
||||
|
||||
async fn fetch(gw: &str, key: Option<&str>, path: &str) -> reqwest::Response {
|
||||
let mut req = reqwest::Client::new().get(format!("{gw}{path}"));
|
||||
if let Some(k) = key {
|
||||
req = req.bearer_auth(k);
|
||||
}
|
||||
req.send().await.unwrap()
|
||||
}
|
||||
|
||||
fn is_mp3(bytes: &[u8]) -> bool {
|
||||
bytes.starts_with(b"ID3") || (bytes.len() > 1 && bytes[0] == 0xFF && bytes[1] & 0xE0 == 0xE0)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn content_transcodes_to_mp3_and_caches_it() {
|
||||
let mock = spawn_mock(Mode::Ok).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let fleet = Arc::new(CortexState::from_config(&config(&mock.url, dir.path())));
|
||||
seed_model(&fleet).await;
|
||||
let gw = serve(&fleet).await;
|
||||
|
||||
let job: Value = submit(&gw, None, song()).await.json().await.unwrap();
|
||||
let id = job["id"].as_str().unwrap().to_string();
|
||||
wait_terminal(&gw, None, &id, 10).await;
|
||||
|
||||
let path = format!("/v1/audio/generations/{id}/content?format=mp3");
|
||||
let resp = fetch(&gw, None, &path).await;
|
||||
assert_eq!(resp.status(), 200);
|
||||
assert_eq!(resp.headers()["content-type"], "audio/mpeg");
|
||||
assert_eq!(
|
||||
resp.headers()["content-disposition"],
|
||||
format!("attachment; filename=\"{id}.mp3\"").as_str()
|
||||
);
|
||||
let mp3 = resp.bytes().await.unwrap();
|
||||
assert!(is_mp3(&mp3), "not an mp3: {:02x?}", &mp3[..8]);
|
||||
let cached = dir.path().join(format!("{id}.mp3"));
|
||||
assert_eq!(std::fs::read(&cached).unwrap(), mp3.as_ref());
|
||||
|
||||
// The second fetch is served from the cache.
|
||||
std::fs::write(&cached, b"cached").unwrap();
|
||||
let again = fetch(&gw, None, &path).await.bytes().await.unwrap();
|
||||
assert_eq!(again.as_ref(), b"cached");
|
||||
|
||||
// The master is untouched and still the default.
|
||||
let flac = fetch(&gw, None, &format!("/v1/audio/generations/{id}/content")).await;
|
||||
assert_eq!(flac.bytes().await.unwrap().as_ref(), FLAC.as_slice());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn response_format_sets_the_default_download() {
|
||||
let mock = spawn_mock(Mode::Ok).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let fleet = Arc::new(CortexState::from_config(&config(&mock.url, dir.path())));
|
||||
seed_model(&fleet).await;
|
||||
let gw = serve(&fleet).await;
|
||||
|
||||
let mut body = song();
|
||||
body["response_format"] = json!("mp3");
|
||||
let job: Value = submit(&gw, None, body).await.json().await.unwrap();
|
||||
let id = job["id"].as_str().unwrap().to_string();
|
||||
let done = wait_terminal(&gw, None, &id, 10).await;
|
||||
// Stored as FLAC whatever the download format.
|
||||
assert_eq!(done["result"]["format"], "flac");
|
||||
assert!(dir.path().join(format!("{id}.flac")).exists());
|
||||
|
||||
let base = format!("/v1/audio/generations/{id}/content");
|
||||
let resp = fetch(&gw, None, &base).await;
|
||||
assert_eq!(resp.headers()["content-type"], "audio/mpeg");
|
||||
// An explicit format wins over the job's default.
|
||||
let resp = fetch(&gw, None, &format!("{base}?format=flac")).await;
|
||||
assert_eq!(resp.headers()["content-type"], "audio/flac");
|
||||
// Unknown formats are refused.
|
||||
for f in ["wav", "ogg"] {
|
||||
let resp = fetch(&gw, None, &format!("{base}?format={f}")).await;
|
||||
assert_eq!(resp.status(), 400, "format {f}");
|
||||
let err: Value = resp.json().await.unwrap();
|
||||
assert_eq!(err["error"]["code"], "invalid_audio_params");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn legacy_wav_results_are_served_as_wav() {
|
||||
let mock = spawn_mock(Mode::LegacyWav).await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let fleet = Arc::new(CortexState::from_config(&config(&mock.url, dir.path())));
|
||||
seed_model(&fleet).await;
|
||||
let gw = serve(&fleet).await;
|
||||
|
||||
let job: Value = submit(&gw, None, song()).await.json().await.unwrap();
|
||||
let id = job["id"].as_str().unwrap().to_string();
|
||||
let done = wait_terminal(&gw, None, &id, 10).await;
|
||||
assert_eq!(done["result"]["format"], "wav");
|
||||
assert!(dir.path().join(format!("{id}.wav")).exists());
|
||||
|
||||
let base = format!("/v1/audio/generations/{id}/content");
|
||||
let resp = fetch(&gw, None, &base).await;
|
||||
assert_eq!(resp.headers()["content-type"], "audio/wav");
|
||||
assert_eq!(resp.bytes().await.unwrap().as_ref(), WAV);
|
||||
let resp = fetch(&gw, None, &format!("{base}?format=mp3")).await;
|
||||
assert_eq!(resp.status(), 409);
|
||||
let err: Value = resp.json().await.unwrap();
|
||||
assert_eq!(err["error"]["code"], "format_unavailable");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -572,7 +706,7 @@ async fn jobs_survive_a_restart() {
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(content.bytes().await.unwrap().as_ref(), WAV);
|
||||
assert_eq!(content.bytes().await.unwrap().as_ref(), FLAC.as_slice());
|
||||
|
||||
// The unfinished one was re-queued and now completes.
|
||||
let resumed = wait_terminal(&gw, Some("key-alice"), &stuck_id, 10).await;
|
||||
@@ -592,13 +726,62 @@ async fn expired_jobs_are_swept() {
|
||||
let job: Value = submit(&gw, None, song()).await.json().await.unwrap();
|
||||
let id = job["id"].as_str().unwrap().to_string();
|
||||
wait_terminal(&gw, None, &id, 10).await;
|
||||
let mp3 = fetch(
|
||||
&gw,
|
||||
None,
|
||||
&format!("/v1/audio/generations/{id}/content?format=mp3"),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(mp3.status(), 200);
|
||||
assert!(dir.path().join(format!("{id}.mp3")).exists());
|
||||
// completed_at has one-second resolution; let it fall behind the cutoff.
|
||||
tokio::time::sleep(Duration::from_millis(1100)).await;
|
||||
|
||||
cortex_gateway::audio_jobs::sweep_expired(&fleet).await;
|
||||
assert_eq!(get_job(&gw, None, &id).await.status(), 404);
|
||||
assert!(!dir.path().join(format!("{id}.json")).exists());
|
||||
assert!(!dir.path().join(format!("{id}.wav")).exists());
|
||||
assert!(!dir.path().join(format!("{id}.flac")).exists());
|
||||
assert!(!dir.path().join(format!("{id}.mp3")).exists());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn submissions_are_refused_when_the_store_is_low_on_disk() {
|
||||
let mock = spawn_mock(Mode::Ok).await;
|
||||
let rates = [("*", 1000)];
|
||||
|
||||
// A floor no filesystem can meet: refused, retryable, nothing written,
|
||||
// nothing reserved.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mut cfg = config(&mock.url, dir.path());
|
||||
cfg.audio_jobs.min_free_mb = u64::MAX;
|
||||
for k in &mut cfg.entitlements.keys {
|
||||
k.hard_cap = Some(1_000_000);
|
||||
}
|
||||
let fleet = Arc::new(CortexState::from_config(&cfg));
|
||||
let mut table = cortex_gateway::rates::RateTable::default();
|
||||
for (model, rate) in rates {
|
||||
table.set(cortex_gateway::rates::AUDIO_SECOND, model, rate);
|
||||
}
|
||||
fleet.rates.replace(table).await;
|
||||
seed_model(&fleet).await;
|
||||
let gw = serve(&fleet).await;
|
||||
|
||||
let resp = submit(&gw, Some("key-alice"), song()).await;
|
||||
assert_eq!(resp.status(), 503);
|
||||
assert_eq!(resp.headers()["retry-after"], "600");
|
||||
let body: Value = resp.json().await.unwrap();
|
||||
assert_eq!(body["error"]["code"], "audio_storage_full");
|
||||
assert_eq!(spent(&fleet, "alice").await, 0);
|
||||
assert_eq!(std::fs::read_dir(dir.path()).unwrap().count(), 0);
|
||||
|
||||
// A floor any filesystem meets: accepted.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mut cfg = config(&mock.url, dir.path());
|
||||
cfg.audio_jobs.min_free_mb = 1;
|
||||
let fleet = Arc::new(CortexState::from_config(&cfg));
|
||||
seed_model(&fleet).await;
|
||||
let gw = serve(&fleet).await;
|
||||
assert_eq!(submit(&gw, None, song()).await.status(), 202);
|
||||
}
|
||||
|
||||
// ── Per-model rates (#386) ───────────────────────────────────────────
|
||||
|
||||
@@ -0,0 +1,308 @@
|
||||
//! Cold-load placement under concurrency.
|
||||
//!
|
||||
//! The mock neuron keeps real VRAM books for one 24 GB card: loads take
|
||||
//! time, `/health` reports what is actually resident, and the card notes
|
||||
//! if it is ever asked to hold more than it has. That last flag is the
|
||||
//! failure being guarded against — two cold-loads deciding from the same
|
||||
//! stale poll, both landing, and every request on the card then failing
|
||||
//! out of memory.
|
||||
|
||||
use axum::extract::{Path, State};
|
||||
use axum::routing::{get, post};
|
||||
use axum::{Json, Router};
|
||||
use cortex_core::config::{
|
||||
EvictionSettings, EvictionStrategy, GatewayConfig, GatewaySettings, NeuronEndpoint,
|
||||
};
|
||||
use cortex_core::discovery::{DeviceInfo, DiscoveryResponse};
|
||||
use cortex_core::node::{ModelEntry, ModelStatus};
|
||||
use cortex_gateway::router::RouteError;
|
||||
use cortex_gateway::state::CortexState;
|
||||
use serde_json::{Value, json};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
const CARD_MB: u64 = 24_000;
|
||||
|
||||
const CATALOGUE: &str = r#"
|
||||
[[models]]
|
||||
id = "music-a"
|
||||
harness = "candle"
|
||||
vram_mb = 14000
|
||||
|
||||
[[models]]
|
||||
id = "music-b"
|
||||
harness = "candle"
|
||||
vram_mb = 12000
|
||||
"#;
|
||||
|
||||
#[derive(Default)]
|
||||
struct Card {
|
||||
resident: HashMap<String, u64>,
|
||||
in_flight: HashMap<String, usize>,
|
||||
overcommitted: bool,
|
||||
unloads: Vec<String>,
|
||||
/// A model the neuron is loading by itself at startup: not resident
|
||||
/// yet, its memory not yet showing.
|
||||
activating: Option<String>,
|
||||
loads: Vec<String>,
|
||||
}
|
||||
|
||||
fn footprint(model: &str) -> u64 {
|
||||
match model {
|
||||
"music-a" => 14_000,
|
||||
"music-b" => 12_000,
|
||||
_ => 1_000,
|
||||
}
|
||||
}
|
||||
|
||||
async fn spawn_neuron(card: Arc<Mutex<Card>>) -> String {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let inference = url.clone();
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/models/load",
|
||||
post(
|
||||
|State(card): State<Arc<Mutex<Card>>>, Json(spec): Json<Value>| async move {
|
||||
let id = spec["model_id"].as_str().unwrap().to_string();
|
||||
// A load takes a while: long enough for a second
|
||||
// request to decide while this one is in progress.
|
||||
tokio::time::sleep(Duration::from_millis(300)).await;
|
||||
let mut card = card.lock().unwrap();
|
||||
card.loads.push(id.clone());
|
||||
card.resident.insert(id.clone(), footprint(&id));
|
||||
if card.resident.values().sum::<u64>() > CARD_MB {
|
||||
card.overcommitted = true;
|
||||
}
|
||||
Json(json!({"status": "loaded"}))
|
||||
},
|
||||
),
|
||||
)
|
||||
.route(
|
||||
"/models/unload",
|
||||
post(
|
||||
|State(card): State<Arc<Mutex<Card>>>, Json(body): Json<Value>| async move {
|
||||
let id = body["model_id"].as_str().unwrap().to_string();
|
||||
let mut card = card.lock().unwrap();
|
||||
card.resident.remove(&id);
|
||||
card.unloads.push(id);
|
||||
Json(json!({"status": "unloaded"}))
|
||||
},
|
||||
),
|
||||
)
|
||||
.route(
|
||||
"/health",
|
||||
get(|State(card): State<Arc<Mutex<Card>>>| async move {
|
||||
let card = card.lock().unwrap();
|
||||
let used: u64 = card.resident.values().sum();
|
||||
let models: Vec<Value> = card
|
||||
.resident
|
||||
.keys()
|
||||
.map(|id| {
|
||||
json!({
|
||||
"id": id,
|
||||
"in_flight": card.in_flight.get(id).copied().unwrap_or(0),
|
||||
"queue_depth": 0
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
let activation = match &card.activating {
|
||||
Some(id) => json!({"state": "pre_warming", "in_progress": id}),
|
||||
None => json!({"state": "ready"}),
|
||||
};
|
||||
Json(json!({
|
||||
"uptime_secs": 1,
|
||||
"activation": activation,
|
||||
"devices": [{
|
||||
"index": 0,
|
||||
"vram_used_mb": used,
|
||||
"vram_free_mb": CARD_MB.saturating_sub(used),
|
||||
"utilization_pct": 0,
|
||||
"temp_c": 40
|
||||
}],
|
||||
"models": models
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/models/{id}/endpoint",
|
||||
get(move |Path(_id): Path<String>| {
|
||||
let url = inference.clone();
|
||||
async move { Json(json!({"url": url})) }
|
||||
}),
|
||||
)
|
||||
.with_state(card);
|
||||
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
|
||||
url
|
||||
}
|
||||
|
||||
/// A gateway with one healthy 24 GB node, its first `/health` already in.
|
||||
async fn gateway(neuron: &str, tag: &str) -> Arc<CortexState> {
|
||||
let path = std::env::temp_dir().join(format!("cortex-cold-load-{tag}.toml"));
|
||||
std::fs::write(&path, CATALOGUE).unwrap();
|
||||
let config = GatewayConfig {
|
||||
gateway: GatewaySettings {
|
||||
listen: "127.0.0.1:0".into(),
|
||||
metrics_listen: "127.0.0.1:0".into(),
|
||||
},
|
||||
eviction: EvictionSettings {
|
||||
strategy: EvictionStrategy::Lru,
|
||||
defrag_after_cycles: 0,
|
||||
},
|
||||
neurons: vec![NeuronEndpoint {
|
||||
name: "gpu".into(),
|
||||
endpoint: neuron.to_string(),
|
||||
}],
|
||||
models_config: path.to_string_lossy().into_owned(),
|
||||
entitlements: Default::default(),
|
||||
upstream: Default::default(),
|
||||
audio_jobs: Default::default(),
|
||||
};
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
{
|
||||
let mut nodes = fleet.nodes.write().await;
|
||||
let node = nodes.get_mut("gpu").unwrap();
|
||||
node.healthy = true;
|
||||
node.discovery = Some(DiscoveryResponse {
|
||||
hostname: "gpu".into(),
|
||||
os: "linux".into(),
|
||||
kernel: "6".into(),
|
||||
cuda_version: None,
|
||||
driver_version: None,
|
||||
devices: vec![DeviceInfo {
|
||||
index: 0,
|
||||
name: "RTX 4090".into(),
|
||||
vram_total_mb: CARD_MB,
|
||||
compute_capability: "8.9".into(),
|
||||
}],
|
||||
harnesses: vec!["candle".into()],
|
||||
cuda_unavailable_reason: None,
|
||||
max_prompt_tokens: 0,
|
||||
});
|
||||
}
|
||||
cortex_gateway::poller::poll_health(&fleet, "gpu", neuron).await;
|
||||
fleet
|
||||
}
|
||||
|
||||
/// Mark `id` resident in the mock and in cortex's view, last used long
|
||||
/// ago — as the poller would have seen it, not as a fresh cold-load.
|
||||
async fn seed_resident(fleet: &CortexState, card: &Mutex<Card>, neuron: &str, id: &str) {
|
||||
card.lock()
|
||||
.unwrap()
|
||||
.resident
|
||||
.insert(id.into(), footprint(id));
|
||||
{
|
||||
let mut nodes = fleet.nodes.write().await;
|
||||
nodes.get_mut("gpu").unwrap().models.insert(
|
||||
id.into(),
|
||||
ModelEntry {
|
||||
id: id.into(),
|
||||
status: ModelStatus::Loaded,
|
||||
last_accessed: Some(chrono::Utc::now() - chrono::Duration::minutes(10)),
|
||||
vram_estimate_mb: Some(footprint(id)),
|
||||
capabilities: Vec::new(),
|
||||
tool_call: false,
|
||||
reasoning: false,
|
||||
limit: None,
|
||||
servable: None,
|
||||
reasoning_budget: Vec::new(),
|
||||
},
|
||||
);
|
||||
}
|
||||
cortex_gateway::poller::poll_health(fleet, "gpu", neuron).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn concurrent_cold_loads_never_overcommit_the_card() {
|
||||
let card = Arc::new(Mutex::new(Card::default()));
|
||||
let neuron = spawn_neuron(Arc::clone(&card)).await;
|
||||
let fleet = gateway(&neuron, "concurrent").await;
|
||||
|
||||
// Both decide while the card is empty; together they do not fit.
|
||||
let (a, b) = tokio::join!(cortex_gateway::router::resolve(&fleet, "music-a"), async {
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
cortex_gateway::router::resolve(&fleet, "music-b").await
|
||||
});
|
||||
|
||||
let card = card.lock().unwrap();
|
||||
assert!(!card.overcommitted, "both models were resident at once");
|
||||
assert!(a.is_ok(), "{a:?}");
|
||||
// music-a was cold-loaded for a request that has yet to reach it, so
|
||||
// music-b may not take its place yet: retry shortly, not a failure.
|
||||
match b {
|
||||
Err(e @ RouteError::NodeBusy { .. }) => {
|
||||
assert_eq!(e.http_status(), 503);
|
||||
assert!(e.retry_after_secs().is_some());
|
||||
}
|
||||
other => panic!("expected NodeBusy, got {other:?}"),
|
||||
}
|
||||
assert_eq!(card.resident.keys().collect::<Vec<_>>(), ["music-a"]);
|
||||
assert!(card.unloads.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_model_with_requests_in_flight_is_not_evicted() {
|
||||
let card = Arc::new(Mutex::new(Card::default()));
|
||||
let neuron = spawn_neuron(Arc::clone(&card)).await;
|
||||
let fleet = gateway(&neuron, "in-flight").await;
|
||||
card.lock().unwrap().in_flight.insert("music-a".into(), 1);
|
||||
seed_resident(&fleet, &card, &neuron, "music-a").await;
|
||||
|
||||
let b = cortex_gateway::router::resolve(&fleet, "music-b").await;
|
||||
|
||||
assert!(matches!(b, Err(RouteError::NodeBusy { .. })), "{b:?}");
|
||||
let card = card.lock().unwrap();
|
||||
assert!(card.unloads.is_empty(), "evicted a busy model");
|
||||
assert!(!card.overcommitted);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn an_idle_model_is_swapped_out() {
|
||||
let card = Arc::new(Mutex::new(Card::default()));
|
||||
let neuron = spawn_neuron(Arc::clone(&card)).await;
|
||||
let fleet = gateway(&neuron, "idle").await;
|
||||
seed_resident(&fleet, &card, &neuron, "music-a").await;
|
||||
|
||||
let b = cortex_gateway::router::resolve(&fleet, "music-b").await;
|
||||
|
||||
assert!(b.is_ok(), "{b:?}");
|
||||
let card = card.lock().unwrap();
|
||||
assert_eq!(card.unloads, ["music-a"]);
|
||||
assert_eq!(card.resident.keys().collect::<Vec<_>>(), ["music-b"]);
|
||||
assert!(!card.overcommitted);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_cold_load_waits_for_the_neuron_s_own_startup_load() {
|
||||
// The neuron has just restarted and is loading its default model.
|
||||
// That load is in no model list yet and its memory has barely
|
||||
// started to fill, so the card still looks empty.
|
||||
let card = Arc::new(Mutex::new(Card {
|
||||
activating: Some("music-a".into()),
|
||||
..Default::default()
|
||||
}));
|
||||
let neuron = spawn_neuron(Arc::clone(&card)).await;
|
||||
let fleet = gateway(&neuron, "activating").await;
|
||||
|
||||
let b = cortex_gateway::router::resolve(&fleet, "music-b").await;
|
||||
|
||||
assert!(matches!(b, Err(RouteError::NodeBusy { .. })), "{b:?}");
|
||||
assert!(
|
||||
card.lock().unwrap().loads.is_empty(),
|
||||
"loaded over a startup load"
|
||||
);
|
||||
|
||||
// Once the startup load finishes and its model is idle, the swap
|
||||
// goes ahead as usual.
|
||||
{
|
||||
let mut c = card.lock().unwrap();
|
||||
c.activating = None;
|
||||
}
|
||||
seed_resident(&fleet, &card, &neuron, "music-a").await;
|
||||
let b = cortex_gateway::router::resolve(&fleet, "music-b").await;
|
||||
assert!(b.is_ok(), "{b:?}");
|
||||
let card = card.lock().unwrap();
|
||||
assert_eq!(card.unloads, ["music-a"]);
|
||||
assert!(!card.overcommitted);
|
||||
}
|
||||
@@ -0,0 +1,154 @@
|
||||
//! `/v1/models` agrees with routing about unhealthy neurons (#376).
|
||||
//!
|
||||
//! When a neuron stops answering, the poller marks it unhealthy but keeps
|
||||
//! its last `/models` and activation snapshot. The router then refuses to
|
||||
//! send it anything, so a model that only that neuron holds cannot be
|
||||
//! served without a cold-load elsewhere. The listing used to report it as
|
||||
//! `loaded: true` anyway — and through tier aliases too — so a client that
|
||||
//! trusted `loaded` picked a model that failed at once.
|
||||
|
||||
mod common;
|
||||
|
||||
use cortex_core::config::{
|
||||
EvictionSettings, EvictionStrategy, GatewayConfig, GatewaySettings, NeuronEndpoint,
|
||||
};
|
||||
use cortex_gateway::state::CortexState;
|
||||
use serde_json::{Value, json};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
fn write_catalogue() -> std::path::PathBuf {
|
||||
static SEQ: AtomicU64 = AtomicU64::new(0);
|
||||
let toml = r#"
|
||||
[[models]]
|
||||
id = "Qwen/large"
|
||||
harness = "candle"
|
||||
min_devices = 1
|
||||
|
||||
[aliases]
|
||||
"helexa/large" = "Qwen/large"
|
||||
"#;
|
||||
let path = std::env::temp_dir().join(format!(
|
||||
"cortex_test_unhealthy_listing.{}.{}.toml",
|
||||
std::process::id(),
|
||||
SEQ.fetch_add(1, Ordering::Relaxed)
|
||||
));
|
||||
std::fs::write(&path, toml).unwrap();
|
||||
path
|
||||
}
|
||||
|
||||
/// Two neurons: `big` holds `Qwen/large` and is pre-warming `Qwen/next`;
|
||||
/// `small` holds `Qwen/small`. Both are polled healthy once.
|
||||
async fn polled_fleet() -> Arc<CortexState> {
|
||||
let big = common::spawn_mock_neuron_with_models_and_health(
|
||||
json!([
|
||||
{"id": "Qwen/large", "harness": "candle", "status": "loaded", "devices": [0], "vram_used_mb": 20000}
|
||||
]),
|
||||
json!({
|
||||
"uptime_secs": 30,
|
||||
"devices": [],
|
||||
"activation": {
|
||||
"state": "pre_warming",
|
||||
"pending": [],
|
||||
"in_progress": "Qwen/next",
|
||||
"completed": [],
|
||||
"failed": []
|
||||
}
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let small = common::spawn_mock_neuron_with_models(json!([
|
||||
{"id": "Qwen/small", "harness": "candle", "status": "loaded", "devices": [0], "vram_used_mb": 4000}
|
||||
]))
|
||||
.await;
|
||||
let config = GatewayConfig {
|
||||
gateway: GatewaySettings {
|
||||
listen: "127.0.0.1:0".into(),
|
||||
metrics_listen: "127.0.0.1:0".into(),
|
||||
},
|
||||
eviction: EvictionSettings {
|
||||
strategy: EvictionStrategy::Lru,
|
||||
defrag_after_cycles: 0,
|
||||
},
|
||||
neurons: vec![
|
||||
NeuronEndpoint {
|
||||
name: "big".into(),
|
||||
endpoint: big,
|
||||
},
|
||||
NeuronEndpoint {
|
||||
name: "small".into(),
|
||||
endpoint: small,
|
||||
},
|
||||
],
|
||||
models_config: write_catalogue().to_string_lossy().into_owned(),
|
||||
..Default::default()
|
||||
};
|
||||
let fleet = Arc::new(CortexState::from_config(&config));
|
||||
cortex_gateway::poller::poll_once(&fleet).await;
|
||||
fleet
|
||||
}
|
||||
|
||||
async fn listing(fleet: Arc<CortexState>) -> Vec<Value> {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let app = cortex_gateway::build_app(fleet);
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
let body: Value = reqwest::get(format!("http://{addr}/v1/models"))
|
||||
.await
|
||||
.unwrap()
|
||||
.json()
|
||||
.await
|
||||
.unwrap();
|
||||
body["data"].as_array().unwrap().clone()
|
||||
}
|
||||
|
||||
fn entry<'a>(data: &'a [Value], id: &str) -> &'a Value {
|
||||
data.iter()
|
||||
.find(|m| m["id"] == id)
|
||||
.unwrap_or_else(|| panic!("{id} should be listed"))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn models_on_an_unhealthy_neuron_are_not_reported_loaded() {
|
||||
let fleet = polled_fleet().await;
|
||||
|
||||
// Control: while both neurons answer, the cached state is live.
|
||||
let data = listing(fleet.clone()).await;
|
||||
assert_eq!(entry(&data, "Qwen/large")["loaded"], true);
|
||||
assert_eq!(entry(&data, "helexa/large")["loaded"], true);
|
||||
assert_eq!(
|
||||
entry(&data, "Qwen/next")["locations"][0]["status"],
|
||||
"loading"
|
||||
);
|
||||
|
||||
// `big` goes dark. The poller flips the flag after three failed
|
||||
// polls (test_poller_marks_unreachable_node_unhealthy) and keeps the
|
||||
// snapshot, which is the state this sets up.
|
||||
fleet.nodes.write().await.get_mut("big").unwrap().healthy = false;
|
||||
|
||||
let data = listing(fleet).await;
|
||||
let large = entry(&data, "Qwen/large");
|
||||
assert_eq!(
|
||||
large["loaded"], false,
|
||||
"a model held only by an unhealthy neuron cannot be served as loaded"
|
||||
);
|
||||
assert_eq!(
|
||||
large["locations"],
|
||||
json!([{"node": "big", "status": "unreachable", "vram_estimate_mb": large["locations"][0]["vram_estimate_mb"]}]),
|
||||
"the location stays visible for diagnostics, marked unreachable"
|
||||
);
|
||||
assert_eq!(
|
||||
entry(&data, "helexa/large")["loaded"],
|
||||
false,
|
||||
"the tier alias must not report its target as loaded either"
|
||||
);
|
||||
assert!(
|
||||
data.iter().all(|m| m["id"] != "Qwen/next"),
|
||||
"a stale activation snapshot must not advertise a model as loading"
|
||||
);
|
||||
|
||||
// The healthy neuron is unaffected.
|
||||
assert_eq!(entry(&data, "Qwen/small")["loaded"], true);
|
||||
}
|
||||
@@ -166,6 +166,11 @@ pub struct ScenarioConfig {
|
||||
/// capability. Defaults to `[1024]`.
|
||||
#[serde(default = "default_image_sizes")]
|
||||
pub image_sizes: Vec<u32>,
|
||||
/// Song lengths (seconds) — one `audio:<secs>` scenario per entry
|
||||
/// (#387), run only against models advertising the `audio`
|
||||
/// capability. Defaults to `[30]`.
|
||||
#[serde(default = "default_audio_durations")]
|
||||
pub audio_durations: Vec<u32>,
|
||||
/// Approximate prompt size (tokens) used by the concurrency scenarios.
|
||||
#[serde(default = "default_concurrency_prompt_tokens")]
|
||||
pub concurrency_prompt_tokens: u32,
|
||||
@@ -202,6 +207,7 @@ impl Default for ScenarioConfig {
|
||||
concurrency_nonstreaming_levels: Vec::new(),
|
||||
concurrency_prompt_tokens: default_concurrency_prompt_tokens(),
|
||||
image_sizes: default_image_sizes(),
|
||||
audio_durations: default_audio_durations(),
|
||||
capability_probes: Vec::new(),
|
||||
}
|
||||
}
|
||||
@@ -280,6 +286,9 @@ fn default_api_listen() -> String {
|
||||
fn default_image_sizes() -> Vec<u32> {
|
||||
vec![1024]
|
||||
}
|
||||
fn default_audio_durations() -> Vec<u32> {
|
||||
vec![30]
|
||||
}
|
||||
|
||||
fn default_prompt_sizes() -> Vec<u32> {
|
||||
vec![128, 4096]
|
||||
|
||||
@@ -268,6 +268,93 @@ impl Scenario for ImageLatencyScenario {
|
||||
}
|
||||
}
|
||||
|
||||
/// Fixed-seed music render (#387): one song of a configured length against
|
||||
/// the neuron's synchronous `/v1/audio/generations`, recording the
|
||||
/// server-side phase timing. The cell dimension is the song length, so the
|
||||
/// real-time factor is `seconds / total_s`. Applies only to models
|
||||
/// advertising the `audio` capability.
|
||||
pub struct AudioLatencyScenario {
|
||||
pub id: String,
|
||||
pub seconds: u32,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Scenario for AudioLatencyScenario {
|
||||
fn id(&self) -> &str {
|
||||
&self.id
|
||||
}
|
||||
|
||||
fn prompt_size(&self) -> u32 {
|
||||
self.seconds
|
||||
}
|
||||
|
||||
fn applies_to(&self, model: &ModelInfo) -> bool {
|
||||
model.capabilities.iter().any(|c| c == "audio")
|
||||
}
|
||||
|
||||
async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics> {
|
||||
let url = ctx
|
||||
.chat_url
|
||||
.replace("/v1/chat/completions", "/v1/audio/generations");
|
||||
let body = serde_json::json!({
|
||||
"model": ctx.model_id,
|
||||
"prompt": "Warm lo-fi hip hop beat with mellow Rhodes chords and dusty vinyl drums, 82 BPM",
|
||||
"duration": self.seconds,
|
||||
"seed": 42,
|
||||
});
|
||||
let start = std::time::Instant::now();
|
||||
let resp = ctx
|
||||
.client
|
||||
.post(&url)
|
||||
.timeout(ctx.timeout)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.context("audio request failed")?;
|
||||
let status = resp.status();
|
||||
let result = resp
|
||||
.headers()
|
||||
.get("x-helexa-audio-result")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(str::to_string);
|
||||
let bytes = resp.bytes().await.context("audio response body")?;
|
||||
anyhow::ensure!(
|
||||
status.is_success(),
|
||||
"audio request {status}: {}",
|
||||
String::from_utf8_lossy(&bytes[..bytes.len().min(512)])
|
||||
);
|
||||
let total_s = start.elapsed().as_secs_f64();
|
||||
// A bench that records latency for broken output is worse than none.
|
||||
// Neurons render FLAC; builds before that rendered WAV.
|
||||
let flac = bytes.len() > 42 && &bytes[..4] == b"fLaC";
|
||||
let wav = bytes.len() > 44 && &bytes[..4] == b"RIFF" && &bytes[8..12] == b"WAVE";
|
||||
anyhow::ensure!(flac || wav, "payload is not a FLAC (or WAV) file");
|
||||
let result: serde_json::Value =
|
||||
serde_json::from_str(&result.context("missing x-helexa-audio-result")?)
|
||||
.context("x-helexa-audio-result is not JSON")?;
|
||||
let timing = &result["usage"]["helexa_timing"];
|
||||
Ok(ScenarioMetrics {
|
||||
ttft_s: timing["plan_ms"].as_u64().unwrap_or(0) as f64 / 1000.0,
|
||||
decode_tps: None,
|
||||
total_s,
|
||||
prompt_tokens: None,
|
||||
completion_tokens: 0,
|
||||
prefill_ms: timing["plan_ms"].as_u64(),
|
||||
decode_ms: timing["generate_ms"].as_u64(),
|
||||
prefill_tokens: None,
|
||||
reasoning_tokens: None,
|
||||
cached_tokens: None,
|
||||
tpot_p95_ms: None,
|
||||
concurrency: None,
|
||||
ttft_p95_s: None,
|
||||
queue_wait_ms_median: None,
|
||||
rejected: None,
|
||||
artifact: None,
|
||||
image_units: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the active scenario set from config: one chat-latency scenario per
|
||||
/// prompt size, plus one concurrency scenario per configured level (#89).
|
||||
/// Concurrency levels default to empty (opt-in), since a burst puts real
|
||||
@@ -308,6 +395,12 @@ pub fn build_scenarios(cfg: &ScenarioConfig) -> Vec<Box<dyn Scenario>> {
|
||||
side_px: side,
|
||||
}) as Box<dyn Scenario>);
|
||||
}
|
||||
for &secs in &cfg.audio_durations {
|
||||
scenarios.push(Box::new(AudioLatencyScenario {
|
||||
id: format!("audio:{secs}"),
|
||||
seconds: secs,
|
||||
}) as Box<dyn Scenario>);
|
||||
}
|
||||
for probe in &cfg.capability_probes {
|
||||
scenarios.push(Box::new(CapabilityScenario {
|
||||
id: format!("capability:{}", probe.name),
|
||||
@@ -916,6 +1009,13 @@ mod tests {
|
||||
assert!(!serves_chat(&model(&["image"])));
|
||||
assert!(!serves_chat(&model(&["decision"])));
|
||||
assert!(!serves_chat(&model(&["audio"])));
|
||||
let audio = AudioLatencyScenario {
|
||||
id: "audio:30".into(),
|
||||
seconds: 30,
|
||||
};
|
||||
assert!(audio.applies_to(&model(&["audio"])));
|
||||
assert!(!audio.applies_to(&model(&["text"])));
|
||||
assert!(!audio.applies_to(&model(&["image"])));
|
||||
let image = ImageLatencyScenario {
|
||||
id: "image:512".into(),
|
||||
side_px: 512,
|
||||
@@ -1070,6 +1170,7 @@ mod tests {
|
||||
use crate::config::{CapabilityProbe, ScenarioConfig};
|
||||
let cfg = ScenarioConfig {
|
||||
image_sizes: vec![],
|
||||
audio_durations: vec![],
|
||||
prompt_sizes: vec![128],
|
||||
max_tokens: 64,
|
||||
concurrency_levels: vec![2, 8],
|
||||
|
||||
@@ -89,6 +89,7 @@ fn config_for(endpoint: String, db_path: String) -> BenchConfig {
|
||||
},
|
||||
scenarios: ScenarioConfig {
|
||||
image_sizes: vec![],
|
||||
audio_durations: vec![],
|
||||
prompt_sizes: vec![128], // single scenario keeps assertions simple
|
||||
max_tokens: 16,
|
||||
concurrency_levels: Vec::new(),
|
||||
|
||||
@@ -14,7 +14,7 @@ use crate::error::envelope_response;
|
||||
use crate::state::RouterState;
|
||||
use axum::Router;
|
||||
use axum::body::{Body, Bytes};
|
||||
use axum::extract::{Path, State};
|
||||
use axum::extract::{Path, RawQuery, State};
|
||||
use axum::http::{HeaderMap, Method, StatusCode};
|
||||
use axum::response::{IntoResponse, Json, Response};
|
||||
use axum::routing::get;
|
||||
@@ -66,15 +66,14 @@ async fn content(
|
||||
State(state): State<Arc<RouterState>>,
|
||||
headers: HeaderMap,
|
||||
Path(id): Path<String>,
|
||||
RawQuery(query): RawQuery,
|
||||
) -> Response {
|
||||
find(
|
||||
&state,
|
||||
Method::GET,
|
||||
&format!("{PATH}/{id}/content"),
|
||||
&headers,
|
||||
&id,
|
||||
)
|
||||
.await
|
||||
// The query carries the download format (`?format=mp3`).
|
||||
let path = match query {
|
||||
Some(q) => format!("{PATH}/{id}/content?{q}"),
|
||||
None => format!("{PATH}/{id}/content"),
|
||||
};
|
||||
find(&state, Method::GET, &path, &headers, &id).await
|
||||
}
|
||||
|
||||
/// Ask each cortex for the job until one answers with something other than
|
||||
|
||||
@@ -170,6 +170,7 @@ mod tests {
|
||||
cost: None,
|
||||
audio: None,
|
||||
audio_tokens_per_second: None,
|
||||
notices: Vec::new(),
|
||||
tool_call: false,
|
||||
reasoning: false,
|
||||
max_model_len: None,
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
//! its cortex, so job calls are asked of each cortex until one answers with
|
||||
//! something other than 404, and listings merge every cortex's jobs.
|
||||
|
||||
use axum::extract::Path;
|
||||
use axum::extract::{Path, RawQuery};
|
||||
use axum::http::{HeaderMap, StatusCode};
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::routing::get;
|
||||
@@ -51,18 +51,25 @@ async fn spawn_cortex(job_id: &'static str, created: u64) -> String {
|
||||
)
|
||||
.route(
|
||||
"/v1/audio/generations/{id}/content",
|
||||
get(move |Path(id): Path<String>| async move {
|
||||
if id == job_id {
|
||||
(
|
||||
StatusCode::OK,
|
||||
[("content-type", "audio/wav")],
|
||||
format!("audio-of-{job_id}"),
|
||||
)
|
||||
.into_response()
|
||||
} else {
|
||||
not_found()
|
||||
}
|
||||
}),
|
||||
get(
|
||||
move |Path(id): Path<String>, RawQuery(q): RawQuery| async move {
|
||||
if id == job_id {
|
||||
// Echo the query so the test can see it was forwarded.
|
||||
let (ct, suffix) = match q.as_deref() {
|
||||
Some("format=mp3") => ("audio/mpeg", ".mp3"),
|
||||
_ => ("audio/flac", ""),
|
||||
};
|
||||
(
|
||||
StatusCode::OK,
|
||||
[("content-type", ct)],
|
||||
format!("audio-of-{job_id}{suffix}"),
|
||||
)
|
||||
.into_response()
|
||||
} else {
|
||||
not_found()
|
||||
}
|
||||
},
|
||||
),
|
||||
);
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
@@ -128,9 +135,20 @@ async fn job_calls_find_the_cortex_that_holds_the_job() {
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(content.status(), 200);
|
||||
assert_eq!(content.headers()["content-type"], "audio/wav");
|
||||
assert_eq!(content.headers()["content-type"], "audio/flac");
|
||||
assert_eq!(content.text().await.unwrap(), "audio-of-audiogen_b");
|
||||
|
||||
// The download format in the query reaches the cortex.
|
||||
let mp3 = http
|
||||
.get(format!(
|
||||
"{router}/v1/audio/generations/audiogen_b/content?format=mp3"
|
||||
))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(mp3.headers()["content-type"], "audio/mpeg");
|
||||
assert_eq!(mp3.text().await.unwrap(), "audio-of-audiogen_b.mp3");
|
||||
|
||||
let cancelled: Value = http
|
||||
.delete(format!("{router}/v1/audio/generations/audiogen_a"))
|
||||
.send()
|
||||
|
||||
@@ -39,6 +39,7 @@ fn model_entry(loaded: bool, feasible: bool) -> CortexModelEntry {
|
||||
cost: None,
|
||||
audio: None,
|
||||
audio_tokens_per_second: None,
|
||||
notices: Vec::new(),
|
||||
tool_call: false,
|
||||
reasoning: false,
|
||||
max_model_len: None,
|
||||
|
||||
@@ -53,6 +53,7 @@ cuda-integration = ["cuda"]
|
||||
|
||||
[dependencies]
|
||||
cortex-core.workspace = true
|
||||
flacenc = { version = "0.5.1", default-features = false, features = ["par"] }
|
||||
tokio.workspace = true
|
||||
axum.workspace = true
|
||||
metrics.workspace = true
|
||||
@@ -147,6 +148,7 @@ base64 = "0.22"
|
||||
tokio = { workspace = true, features = ["test-util"] }
|
||||
reqwest.workspace = true
|
||||
tempfile = "3"
|
||||
claxon = "0.4.3"
|
||||
|
||||
[build-dependencies]
|
||||
# Used by `build.rs` to compile `src/cuda/*.cu` into `libneuroncuda.a`
|
||||
|
||||
@@ -142,8 +142,22 @@ async fn load_model(
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
let registry = state.registry.read().await;
|
||||
match registry.load_model(&spec).await {
|
||||
// Run the load on its own task. A client that gives up (a curl or
|
||||
// cortex timeout during a first fetch of tens of GB) drops this
|
||||
// handler's future; the load must not be dropped with it, or the
|
||||
// download restarts from the last finished file and a device-worker
|
||||
// reply arrives with nobody to receive it.
|
||||
let load_spec = spec.clone();
|
||||
let load_state = Arc::clone(&state);
|
||||
let load = tokio::spawn(async move {
|
||||
let registry = load_state.registry.read().await;
|
||||
registry.load_model(&load_spec).await
|
||||
});
|
||||
let result = match load.await {
|
||||
Ok(r) => r,
|
||||
Err(e) => Err(anyhow::anyhow!("load task failed: {e}")),
|
||||
};
|
||||
match result {
|
||||
Ok(()) => Json(json!({"status": "loaded"})).into_response(),
|
||||
Err(e) => {
|
||||
// If the underlying failure is a structured preflight
|
||||
@@ -674,7 +688,7 @@ async fn images_generations(
|
||||
/// `POST /v1/audio/generations` (#382): render one audio generation and
|
||||
/// return it. This is the neuron half of the async job API — cortex owns
|
||||
/// the job and calls this once per job, waiting for the whole render. The
|
||||
/// body is the audio (`audio/wav`); the result metadata travels in the
|
||||
/// body is the audio (`audio/flac`); the result metadata travels in the
|
||||
/// [`cortex_core::audio::RESULT_HEADER`] header so the bytes need no
|
||||
/// base64 envelope.
|
||||
async fn audio_generations(
|
||||
@@ -702,7 +716,7 @@ async fn audio_generations(
|
||||
}
|
||||
let principal = principal_key(&headers);
|
||||
match candle.audio_generation(&req, principal.as_deref()).await {
|
||||
Ok((wav, result)) => {
|
||||
Ok((flac, result)) => {
|
||||
let meta = match serde_json::to_string(&result) {
|
||||
Ok(m) => m,
|
||||
Err(e) => {
|
||||
@@ -711,11 +725,13 @@ async fn audio_generations(
|
||||
)));
|
||||
}
|
||||
};
|
||||
let mut resp = axum::response::Response::new(axum::body::Body::from(wav));
|
||||
let mut resp = axum::response::Response::new(axum::body::Body::from(flac));
|
||||
let h = resp.headers_mut();
|
||||
h.insert(
|
||||
axum::http::header::CONTENT_TYPE,
|
||||
axum::http::HeaderValue::from_static("audio/wav"),
|
||||
axum::http::HeaderValue::from_static(cortex_core::audio::content_type(
|
||||
cortex_core::audio::STORED_FORMAT,
|
||||
)),
|
||||
);
|
||||
if let Ok(v) = axum::http::HeaderValue::from_str(&meta) {
|
||||
h.insert(RESULT_HEADER, v);
|
||||
|
||||
@@ -138,6 +138,66 @@ pub struct CandleHarnessConfig {
|
||||
/// text-to-image pipeline.
|
||||
#[serde(default)]
|
||||
pub image: ImageConfig,
|
||||
|
||||
/// `[harness.candle.audio]` settings (#385) — how an ACE-Step
|
||||
/// pipeline is laid out on this host.
|
||||
#[serde(default)]
|
||||
pub audio: AudioConfig,
|
||||
}
|
||||
|
||||
/// `[harness.candle.audio]` settings (#385).
|
||||
///
|
||||
/// These describe the host, not the model, which is why they live here
|
||||
/// and not in the catalogue: the same ACE-Step id serves bf16 on a 4090
|
||||
/// and q8 with f32 compute on an 8 GB Pascal card.
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct AudioConfig {
|
||||
/// In-situ weight quantization of the DiT, the text encoder and the
|
||||
/// planner: `"q8_0"`, or unset for none. Quantized components compute
|
||||
/// in f32. The VAE is never quantized.
|
||||
#[serde(default)]
|
||||
pub quant: Option<String>,
|
||||
/// Compute dtype of the unquantized GPU components: `"bf16"`
|
||||
/// (default) or `"f32"`. Cards without fast bf16 (Pascal) need f32.
|
||||
#[serde(default)]
|
||||
pub compute: Option<String>,
|
||||
/// Where each component runs.
|
||||
#[serde(default)]
|
||||
pub placement: AudioPlacement,
|
||||
/// Which planner LM to load: a snapshot subfolder (for example
|
||||
/// `"acestep-5Hz-lm-1.7B"`), a repo id of its own (the 0.6B planner is
|
||||
/// `"ACE-Step/acestep-5Hz-lm-0.6B"`), or `"none"`. Unset takes the
|
||||
/// largest the snapshot has.
|
||||
#[serde(default)]
|
||||
pub planner: Option<String>,
|
||||
/// Treat the device as having this much free memory for the pipeline,
|
||||
/// in MiB. Under a budget the VAE decodes in smaller chunks.
|
||||
#[serde(default)]
|
||||
pub vram_budget_mb: Option<u64>,
|
||||
/// Make the budget real by holding the rest of the device's free
|
||||
/// memory while the pipeline is loaded, so a larger card fails exactly
|
||||
/// where a card of the budget's size would. For rehearsing a small
|
||||
/// card's profile on a large one; never for serving.
|
||||
#[serde(default)]
|
||||
pub enforce_vram_budget: bool,
|
||||
/// Longest song this host accepts, in seconds, when it is less than
|
||||
/// the model's own limit.
|
||||
#[serde(default)]
|
||||
pub max_duration_secs: Option<f64>,
|
||||
}
|
||||
|
||||
/// Per-component placement (#385): `"gpu"` (the model's device, the
|
||||
/// default) or `"cpu"`. CPU components compute in f32.
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct AudioPlacement {
|
||||
#[serde(default)]
|
||||
pub text_encoder: Option<String>,
|
||||
#[serde(default)]
|
||||
pub planner: Option<String>,
|
||||
#[serde(default)]
|
||||
pub dit: Option<String>,
|
||||
#[serde(default)]
|
||||
pub vae: Option<String>,
|
||||
}
|
||||
|
||||
/// `[harness.candle.image]` settings (#198).
|
||||
|
||||
@@ -13,9 +13,10 @@
|
||||
//! All encoder layers share the same architecture: pre-norm (RmsNorm) bidirectional
|
||||
//! GQA attention followed by a gated SiLU MLP, matching the Qwen3 block design.
|
||||
|
||||
use super::linear::{Linear, linear_b, linear_no_bias};
|
||||
use candle_core::{DType, Device, Module, Result, Tensor};
|
||||
use candle_nn::{Activation, VarBuilder};
|
||||
use candle_transformers::models::with_tracing::{Linear, RmsNorm, linear_b, linear_no_bias};
|
||||
use candle_transformers::models::with_tracing::RmsNorm;
|
||||
use candle_transformers::utils::repeat_kv;
|
||||
|
||||
use super::AceStepConfig;
|
||||
|
||||
@@ -0,0 +1,615 @@
|
||||
//! The planner's chain of thought, as text (#384).
|
||||
//!
|
||||
//! The planner writes its plan as YAML between `<think>` tags. Upstream
|
||||
//! reads it back with `parse_lm_output`, then writes it out again with
|
||||
//! `yaml.dump` to build the prompt the audio codes are generated from. Those
|
||||
//! bytes are model input, so the dump has to be reproduced exactly,
|
||||
//! including where PyYAML folds a long caption and when it quotes a value.
|
||||
//!
|
||||
//! Reference: `acestep/llm_inference.py` (`parse_lm_output`,
|
||||
//! `_format_metadata_as_cot`) and `acestep/constrained_logits_processor.py`
|
||||
//! (`postprocess_caption`) at ACE-Step 1.5 `ca1e85fe`; PyYAML's
|
||||
//! `emitter.py` for the scalar writers.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
/// A plan field's value as upstream holds it: `int(...)` where the text is
|
||||
/// a whole number, the text otherwise. The distinction decides how YAML
|
||||
/// writes it.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum MetaValue {
|
||||
Int(i64),
|
||||
Str(String),
|
||||
}
|
||||
|
||||
impl MetaValue {
|
||||
pub fn as_text(&self) -> String {
|
||||
match self {
|
||||
Self::Int(i) => i.to_string(),
|
||||
Self::Str(s) => s.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The plan's fields, by name.
|
||||
pub type Metadata = BTreeMap<String, MetaValue>;
|
||||
|
||||
/// Fields the codes-phase prompt carries, in the order YAML sorts them.
|
||||
const COT_FIELDS: [&str; 6] = [
|
||||
"bpm",
|
||||
"caption",
|
||||
"duration",
|
||||
"keyscale",
|
||||
"language",
|
||||
"timesignature",
|
||||
];
|
||||
|
||||
/// Join a YAML-folded caption back into one line: every line trimmed,
|
||||
/// blank lines dropped, joined with single spaces.
|
||||
pub fn postprocess_caption(caption: &str) -> String {
|
||||
caption
|
||||
.split('\n')
|
||||
.map(str::trim)
|
||||
.filter(|l| !l.is_empty())
|
||||
.collect::<Vec<_>>()
|
||||
.join(" ")
|
||||
}
|
||||
|
||||
/// Python's `str.isdigit` for the ASCII digits upstream's values use.
|
||||
fn is_digits(s: &str) -> bool {
|
||||
!s.is_empty() && s.chars().all(|c| c.is_ascii_digit())
|
||||
}
|
||||
|
||||
/// Read a plan out of the planner's output, as `parse_lm_output` does.
|
||||
///
|
||||
/// The CoT phase stops before `</think>`, so the reasoning section is the
|
||||
/// text between the tags when both are present, else everything before the
|
||||
/// first audio code. A field is a line that does not start with whitespace
|
||||
/// and has a colon; indented lines continue the previous field.
|
||||
pub fn parse_plan(output: &str) -> Metadata {
|
||||
let reasoning = match (output.find("<think>"), output.find("</think>")) {
|
||||
(Some(s), Some(e)) if e > s => output[s + "<think>".len()..e].trim().to_string(),
|
||||
_ => output
|
||||
.split("<|audio_code_")
|
||||
.next()
|
||||
.unwrap_or("")
|
||||
.trim()
|
||||
.to_string(),
|
||||
};
|
||||
|
||||
let mut meta = Metadata::new();
|
||||
let mut key: Option<String> = None;
|
||||
let mut lines: Vec<&str> = Vec::new();
|
||||
let save = |key: &mut Option<String>, lines: &mut Vec<&str>, meta: &mut Metadata| {
|
||||
if let Some(k) = key.take()
|
||||
&& !lines.is_empty()
|
||||
{
|
||||
let value = lines.join("\n");
|
||||
let v = match k.as_str() {
|
||||
"bpm" | "duration" => match value.trim().parse::<i64>() {
|
||||
Ok(i) => MetaValue::Int(i),
|
||||
Err(_) => MetaValue::Str(value.trim().to_string()),
|
||||
},
|
||||
"caption" => MetaValue::Str(postprocess_caption(&value)),
|
||||
"genres" | "keyscale" | "language" | "timesignature" => {
|
||||
MetaValue::Str(value.trim().to_string())
|
||||
}
|
||||
_ => {
|
||||
lines.clear();
|
||||
return;
|
||||
}
|
||||
};
|
||||
meta.insert(k, v);
|
||||
}
|
||||
lines.clear();
|
||||
};
|
||||
for line in reasoning.split('\n') {
|
||||
if line.trim().starts_with('<') {
|
||||
continue;
|
||||
}
|
||||
let starts_blank = line.chars().next().is_some_and(char::is_whitespace);
|
||||
if !line.is_empty() && !starts_blank && line.contains(':') {
|
||||
save(&mut key, &mut lines, &mut meta);
|
||||
let (k, v) = line.split_once(':').unwrap_or((line, ""));
|
||||
key = Some(k.trim().to_lowercase());
|
||||
if !v.trim().is_empty() {
|
||||
lines.push(v);
|
||||
}
|
||||
} else if (line.starts_with(' ') || line.starts_with('\t')) && key.is_some() {
|
||||
lines.push(line);
|
||||
}
|
||||
}
|
||||
save(&mut key, &mut lines, &mut meta);
|
||||
meta
|
||||
}
|
||||
|
||||
/// The CoT block the codes phase is prompted with, as
|
||||
/// `_format_metadata_as_cot` builds it: `<think>\n{yaml}\n</think>`.
|
||||
pub fn format_plan(meta: &Metadata) -> String {
|
||||
let mut yaml = String::new();
|
||||
for key in COT_FIELDS {
|
||||
let Some(value) = meta.get(key) else {
|
||||
continue;
|
||||
};
|
||||
let value = match value {
|
||||
MetaValue::Str(s) if key == "timesignature" && s.ends_with("/4") => {
|
||||
MetaValue::Str(s.split('/').next().unwrap_or("").to_string())
|
||||
}
|
||||
v => v.clone(),
|
||||
};
|
||||
let value = match value {
|
||||
MetaValue::Str(s) if is_digits(&s) => match s.parse::<i64>() {
|
||||
Ok(i) => MetaValue::Int(i),
|
||||
Err(_) => MetaValue::Str(s),
|
||||
},
|
||||
v => v,
|
||||
};
|
||||
yaml.push_str(&yaml_entry(key, &value));
|
||||
}
|
||||
format!("<think>\n{}\n</think>", yaml.trim())
|
||||
}
|
||||
|
||||
// ── A PyYAML block-mapping entry ─────────────────────────────────────
|
||||
|
||||
/// PyYAML's default line width and mapping indent.
|
||||
const BEST_WIDTH: usize = 80;
|
||||
const INDENT: usize = 2;
|
||||
|
||||
/// One `key: value\n` line (or lines) of `yaml.dump(..., allow_unicode=True)`.
|
||||
fn yaml_entry(key: &str, value: &MetaValue) -> String {
|
||||
let mut w = Writer {
|
||||
out: format!("{key}:"),
|
||||
column: key.chars().count() + 1,
|
||||
};
|
||||
match value {
|
||||
MetaValue::Int(i) => {
|
||||
w.out.push(' ');
|
||||
w.out.push_str(&i.to_string());
|
||||
}
|
||||
MetaValue::Str(s) => match scalar_style(s) {
|
||||
Style::Plain => w.plain(s),
|
||||
Style::SingleQuoted => w.single_quoted(s),
|
||||
Style::DoubleQuoted => w.double_quoted(s),
|
||||
},
|
||||
}
|
||||
w.out.push('\n');
|
||||
w.out
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq)]
|
||||
enum Style {
|
||||
Plain,
|
||||
SingleQuoted,
|
||||
DoubleQuoted,
|
||||
}
|
||||
|
||||
fn is_ws(c: Option<char>) -> bool {
|
||||
matches!(
|
||||
c,
|
||||
None | Some('\0' | ' ' | '\t' | '\r' | '\n' | '\u{85}' | '\u{2028}' | '\u{2029}')
|
||||
)
|
||||
}
|
||||
|
||||
fn is_break(c: char) -> bool {
|
||||
matches!(c, '\n' | '\u{85}' | '\u{2028}' | '\u{2029}')
|
||||
}
|
||||
|
||||
/// `choose_scalar_style` for a block-mapping value: plain when PyYAML's
|
||||
/// analysis allows a block plain scalar and the text would read back as a
|
||||
/// string, else single-quoted when allowed, else double-quoted.
|
||||
fn scalar_style(s: &str) -> Style {
|
||||
if s.is_empty() {
|
||||
// An empty value would read back as null.
|
||||
return Style::SingleQuoted;
|
||||
}
|
||||
let chars: Vec<char> = s.chars().collect();
|
||||
let n = chars.len();
|
||||
let mut block_indicators = s.starts_with("---") || s.starts_with("...");
|
||||
let mut line_breaks = false;
|
||||
let mut special = false;
|
||||
let (mut leading_space, mut leading_break, mut trailing_space, mut trailing_break) =
|
||||
(false, false, false, false);
|
||||
let (mut break_space, mut space_break) = (false, false);
|
||||
let (mut previous_space, mut previous_break) = (false, false);
|
||||
let mut preceded_by_ws = true;
|
||||
for i in 0..n {
|
||||
let ch = chars[i];
|
||||
let followed_by_ws = is_ws(chars.get(i + 1).copied());
|
||||
if i == 0 {
|
||||
if "#,[]{}&*!|>'\"%@`".contains(ch) {
|
||||
block_indicators = true;
|
||||
}
|
||||
if "?:".contains(ch) && followed_by_ws {
|
||||
block_indicators = true;
|
||||
}
|
||||
if ch == '-' && followed_by_ws {
|
||||
block_indicators = true;
|
||||
}
|
||||
} else {
|
||||
if ch == ':' && followed_by_ws {
|
||||
block_indicators = true;
|
||||
}
|
||||
if ch == '#' && preceded_by_ws {
|
||||
block_indicators = true;
|
||||
}
|
||||
}
|
||||
if is_break(ch) {
|
||||
line_breaks = true;
|
||||
}
|
||||
if !(ch == '\n' || (' '..='~').contains(&ch)) {
|
||||
let printable = ch == '\u{85}'
|
||||
|| ('\u{A0}'..='\u{D7FF}').contains(&ch)
|
||||
|| ('\u{E000}'..='\u{FFFD}').contains(&ch)
|
||||
|| ('\u{10000}'..'\u{10FFFF}').contains(&ch);
|
||||
// allow_unicode=True: printable non-ASCII is fine.
|
||||
if !printable || ch == '\u{FEFF}' {
|
||||
special = true;
|
||||
}
|
||||
}
|
||||
if ch == ' ' {
|
||||
leading_space |= i == 0;
|
||||
trailing_space |= i == n - 1;
|
||||
break_space |= previous_break;
|
||||
previous_space = true;
|
||||
previous_break = false;
|
||||
} else if is_break(ch) {
|
||||
leading_break |= i == 0;
|
||||
trailing_break |= i == n - 1;
|
||||
space_break |= previous_space;
|
||||
previous_space = false;
|
||||
previous_break = true;
|
||||
} else {
|
||||
previous_space = false;
|
||||
previous_break = false;
|
||||
}
|
||||
preceded_by_ws = is_ws(Some(ch));
|
||||
}
|
||||
let mut allow_plain = true;
|
||||
let mut allow_single = true;
|
||||
if leading_space || leading_break || trailing_space || trailing_break {
|
||||
allow_plain = false;
|
||||
}
|
||||
if break_space {
|
||||
allow_plain = false;
|
||||
allow_single = false;
|
||||
}
|
||||
if space_break || special {
|
||||
allow_plain = false;
|
||||
allow_single = false;
|
||||
}
|
||||
if line_breaks || block_indicators {
|
||||
allow_plain = false;
|
||||
}
|
||||
if allow_plain && resolves_as_string(s) {
|
||||
Style::Plain
|
||||
} else if allow_single {
|
||||
Style::SingleQuoted
|
||||
} else {
|
||||
Style::DoubleQuoted
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether a plain scalar with this text would read back as a string under
|
||||
/// PyYAML's YAML 1.1 implicit resolvers (bool, int, float, null, merge,
|
||||
/// value, timestamp). Anything else needs quoting.
|
||||
fn resolves_as_string(s: &str) -> bool {
|
||||
const BOOLS: [&str; 18] = [
|
||||
"yes", "Yes", "YES", "no", "No", "NO", "true", "True", "TRUE", "false", "False", "FALSE",
|
||||
"on", "On", "ON", "off", "Off", "OFF",
|
||||
];
|
||||
const NULLS: [&str; 4] = ["~", "null", "Null", "NULL"];
|
||||
if BOOLS.contains(&s) || NULLS.contains(&s) || s == "<<" || s == "=" {
|
||||
return false;
|
||||
}
|
||||
!(is_yaml_int(s) || is_yaml_float(s) || is_yaml_timestamp(s))
|
||||
}
|
||||
|
||||
fn strip_sign(s: &str) -> &str {
|
||||
s.strip_prefix(['-', '+']).unwrap_or(s)
|
||||
}
|
||||
|
||||
/// PyYAML's int resolver: binary, octal, decimal, hex and sexagesimal,
|
||||
/// with `_` separators.
|
||||
fn is_yaml_int(s: &str) -> bool {
|
||||
let b = strip_sign(s);
|
||||
let all = |t: &str, ok: fn(char) -> bool| !t.is_empty() && t.chars().all(|c| ok(c) || c == '_');
|
||||
if let Some(t) = b.strip_prefix("0b") {
|
||||
return all(t, |c| c == '0' || c == '1');
|
||||
}
|
||||
if let Some(t) = b.strip_prefix("0x") {
|
||||
return all(t, |c| c.is_ascii_hexdigit());
|
||||
}
|
||||
if b == "0" {
|
||||
return true;
|
||||
}
|
||||
if let Some(t) = b.strip_prefix('0') {
|
||||
return all(t, |c| ('0'..='7').contains(&c));
|
||||
}
|
||||
if b.starts_with(|c: char| ('1'..='9').contains(&c)) {
|
||||
if all(b, |c| c.is_ascii_digit()) {
|
||||
return true;
|
||||
}
|
||||
// Sexagesimal: [1-9][0-9_]*(:[0-5]?[0-9])+
|
||||
let mut parts = b.split(':');
|
||||
let head = parts.next().unwrap_or("");
|
||||
let rest: Vec<&str> = parts.collect();
|
||||
return !rest.is_empty()
|
||||
&& all(head, |c| c.is_ascii_digit())
|
||||
&& rest.iter().all(|p| {
|
||||
let p: Vec<char> = p.chars().collect();
|
||||
match p.as_slice() {
|
||||
[d] => d.is_ascii_digit(),
|
||||
[a, d] => ('0'..='5').contains(a) && d.is_ascii_digit(),
|
||||
_ => false,
|
||||
}
|
||||
});
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// PyYAML's float resolver, including `.inf`, `.nan` and sexagesimal.
|
||||
fn is_yaml_float(s: &str) -> bool {
|
||||
let b = strip_sign(s);
|
||||
if matches!(b, ".inf" | ".Inf" | ".INF") {
|
||||
return true;
|
||||
}
|
||||
if matches!(s, ".nan" | ".NaN" | ".NAN") {
|
||||
return true;
|
||||
}
|
||||
// [-+]?(?:[0-9][0-9_]*)\.[0-9_]*(?:[eE][-+][0-9]+)? and \.[0-9][0-9_]*(?:[eE][-+][0-9]+)?
|
||||
let (mantissa, exp) = match b.find(['e', 'E']) {
|
||||
Some(i) => (&b[..i], Some(&b[i + 1..])),
|
||||
None => (b, None),
|
||||
};
|
||||
if let Some(e) = exp {
|
||||
let digits = e.strip_prefix(['-', '+']);
|
||||
if !digits.is_some_and(|d| !d.is_empty() && d.chars().all(|c| c.is_ascii_digit())) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
let Some((int, frac)) = mantissa.split_once('.') else {
|
||||
// Sexagesimal float: [0-9][0-9_]*(:[0-5]?[0-9])+\.[0-9_]*
|
||||
return false;
|
||||
};
|
||||
let digits_ = |t: &str| t.chars().all(|c| c.is_ascii_digit() || c == '_');
|
||||
if int.is_empty() {
|
||||
frac.starts_with(|c: char| c.is_ascii_digit()) && digits_(frac)
|
||||
} else {
|
||||
int.starts_with(|c: char| c.is_ascii_digit()) && digits_(int) && digits_(frac)
|
||||
}
|
||||
}
|
||||
|
||||
/// PyYAML's timestamp resolver, enough to catch dates and date-times.
|
||||
fn is_yaml_timestamp(s: &str) -> bool {
|
||||
let c: Vec<char> = s.chars().collect();
|
||||
c.len() >= 8
|
||||
&& c[..4].iter().all(char::is_ascii_digit)
|
||||
&& c[4] == '-'
|
||||
&& c[5].is_ascii_digit()
|
||||
&& c.iter().skip(5).take_while(|x| x.is_ascii_digit()).count() <= 2
|
||||
&& s[5..].contains('-')
|
||||
}
|
||||
|
||||
/// The emitter state a scalar writer needs: the text so far and the
|
||||
/// current column.
|
||||
struct Writer {
|
||||
out: String,
|
||||
column: usize,
|
||||
}
|
||||
|
||||
impl Writer {
|
||||
fn write(&mut self, s: &str) {
|
||||
self.column += s.chars().count();
|
||||
self.out.push_str(s);
|
||||
}
|
||||
|
||||
/// `write_indent` mid-scalar: a line break and the mapping indent.
|
||||
fn indent(&mut self) {
|
||||
self.out.push('\n');
|
||||
self.out.push_str(&" ".repeat(INDENT));
|
||||
self.column = INDENT;
|
||||
}
|
||||
|
||||
/// `write_plain` with `split=True`: a run of exactly one space past
|
||||
/// the best width becomes a line break.
|
||||
fn plain(&mut self, text: &str) {
|
||||
self.write(" ");
|
||||
let chars: Vec<char> = text.chars().collect();
|
||||
let (mut spaces, mut start, mut end) = (false, 0usize, 0usize);
|
||||
while end <= chars.len() {
|
||||
let ch = chars.get(end).copied();
|
||||
if spaces {
|
||||
if ch != Some(' ') {
|
||||
if start + 1 == end && self.column > BEST_WIDTH {
|
||||
self.indent();
|
||||
} else {
|
||||
self.write(&chars[start..end].iter().collect::<String>());
|
||||
}
|
||||
start = end;
|
||||
}
|
||||
} else if ch.is_none() || ch == Some(' ') {
|
||||
self.write(&chars[start..end].iter().collect::<String>());
|
||||
start = end;
|
||||
}
|
||||
if let Some(c) = ch {
|
||||
spaces = c == ' ';
|
||||
}
|
||||
end += 1;
|
||||
}
|
||||
}
|
||||
|
||||
/// `write_single_quoted` with `split=True` (no line breaks reach it:
|
||||
/// captions are joined to one line first).
|
||||
fn single_quoted(&mut self, text: &str) {
|
||||
self.write(" '");
|
||||
let chars: Vec<char> = text.chars().collect();
|
||||
let n = chars.len();
|
||||
let (mut spaces, mut start, mut end) = (false, 0usize, 0usize);
|
||||
while end <= n {
|
||||
let ch = chars.get(end).copied();
|
||||
if spaces {
|
||||
if ch != Some(' ') {
|
||||
if start + 1 == end && self.column > BEST_WIDTH && start != 0 && end != n {
|
||||
self.indent();
|
||||
} else {
|
||||
self.write(&chars[start..end].iter().collect::<String>());
|
||||
}
|
||||
start = end;
|
||||
}
|
||||
} else if (ch.is_none() || ch == Some(' ') || ch == Some('\'')) && start < end {
|
||||
self.write(&chars[start..end].iter().collect::<String>());
|
||||
start = end;
|
||||
}
|
||||
if ch == Some('\'') {
|
||||
self.write("''");
|
||||
start = end + 1;
|
||||
}
|
||||
if let Some(c) = ch {
|
||||
spaces = c == ' ';
|
||||
}
|
||||
end += 1;
|
||||
}
|
||||
self.write("'");
|
||||
}
|
||||
|
||||
/// Double quotes for text with control characters. PyYAML folds these
|
||||
/// too; a plan never contains them, so this writes one line.
|
||||
fn double_quoted(&mut self, text: &str) {
|
||||
let mut s = String::from(" \"");
|
||||
for c in text.chars() {
|
||||
match c {
|
||||
'"' => s.push_str("\\\""),
|
||||
'\\' => s.push_str("\\\\"),
|
||||
'\n' => s.push_str("\\n"),
|
||||
'\t' => s.push_str("\\t"),
|
||||
c if (c as u32) < 0x20 => s.push_str(&format!("\\x{:02X}", c as u32)),
|
||||
c => s.push(c),
|
||||
}
|
||||
}
|
||||
s.push('"');
|
||||
self.write(&s);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn fixture_planner() -> serde_json::Value {
|
||||
let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
|
||||
.join("src/harness/testdata/acestep/vocal_en_30s_plan__cpu_f32.json");
|
||||
let meta: serde_json::Value =
|
||||
serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap();
|
||||
meta["planner"].clone()
|
||||
}
|
||||
|
||||
/// The codes-phase prompt upstream built from its own CoT output: the
|
||||
/// plan read back, then dumped as YAML, folded at column 80.
|
||||
#[test]
|
||||
fn the_plan_round_trips_to_upstream_s_codes_prompt() {
|
||||
let planner = fixture_planner();
|
||||
let cot_out = planner[0]["output_text"].as_str().unwrap();
|
||||
let codes_prompt = planner[1]["formatted_prompt"].as_str().unwrap();
|
||||
let plan = parse_plan(cot_out);
|
||||
assert_eq!(plan["bpm"], MetaValue::Int(130));
|
||||
assert_eq!(plan["duration"], MetaValue::Int(30));
|
||||
assert_eq!(plan["timesignature"], MetaValue::Str("4".into()));
|
||||
let MetaValue::Str(caption) = &plan["caption"] else {
|
||||
panic!("caption is text");
|
||||
};
|
||||
assert!(caption.starts_with("An energetic indie rock track"));
|
||||
assert!(!caption.contains('\n'), "the caption is unfolded");
|
||||
|
||||
let cot = format_plan(&plan);
|
||||
let start = codes_prompt.find("<think>").unwrap();
|
||||
let want = &codes_prompt[start..codes_prompt.len() - "\n\n".len()];
|
||||
assert_eq!(cot, want);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn values_that_would_read_back_as_something_else_are_quoted() {
|
||||
assert_eq!(
|
||||
yaml_entry("language", &MetaValue::Str("no".into())),
|
||||
"language: 'no'\n"
|
||||
);
|
||||
assert_eq!(
|
||||
yaml_entry("language", &MetaValue::Str("en".into())),
|
||||
"language: en\n"
|
||||
);
|
||||
assert_eq!(
|
||||
yaml_entry("keyscale", &MetaValue::Str("F# minor".into())),
|
||||
"keyscale: F# minor\n"
|
||||
);
|
||||
assert_eq!(
|
||||
yaml_entry("caption", &MetaValue::Str("Lo-fi: dusty drums".into())),
|
||||
"caption: 'Lo-fi: dusty drums'\n"
|
||||
);
|
||||
assert_eq!(
|
||||
yaml_entry("caption", &MetaValue::Str("'quoted' start".into())),
|
||||
"caption: '''quoted'' start'\n"
|
||||
);
|
||||
assert_eq!(
|
||||
yaml_entry("caption", &MetaValue::Str("1.5".into())),
|
||||
"caption: '1.5'\n"
|
||||
);
|
||||
assert_eq!(yaml_entry("bpm", &MetaValue::Int(92)), "bpm: 92\n");
|
||||
}
|
||||
|
||||
/// Recorded from PyYAML 6 (`yaml.dump({"caption": v}, allow_unicode=True,
|
||||
/// sort_keys=True)`).
|
||||
#[test]
|
||||
fn scalars_match_pyyaml() {
|
||||
let folded = "A track: ".to_string() + &"word ".repeat(30) + "end";
|
||||
let cases: Vec<(String, &str)> = vec![
|
||||
("no".into(), "caption: 'no'\n"),
|
||||
("en".into(), "caption: en\n"),
|
||||
("F# minor".into(), "caption: F# minor\n"),
|
||||
(
|
||||
"Lo-fi: dusty drums".into(),
|
||||
"caption: 'Lo-fi: dusty drums'\n",
|
||||
),
|
||||
("'quoted' start".into(), "caption: '''quoted'' start'\n"),
|
||||
("1.5".into(), "caption: '1.5'\n"),
|
||||
(
|
||||
folded,
|
||||
"caption: 'A track: word word word word word word word word word word word word word\n word word word word word word word word word word word word word word word word\n word end'\n",
|
||||
),
|
||||
("yes".into(), "caption: 'yes'\n"),
|
||||
("2024-01-02".into(), "caption: '2024-01-02'\n"),
|
||||
("1:30".into(), "caption: '1:30'\n"),
|
||||
("0x1F".into(), "caption: '0x1F'\n"),
|
||||
("~".into(), "caption: '~'\n"),
|
||||
("null".into(), "caption: 'null'\n"),
|
||||
("- dash start".into(), "caption: '- dash start'\n"),
|
||||
("tag #x".into(), "caption: 'tag #x'\n"),
|
||||
("a#b".into(), "caption: a#b\n"),
|
||||
("ends:".into(), "caption: 'ends:'\n"),
|
||||
];
|
||||
for (v, want) in cases {
|
||||
assert_eq!(
|
||||
yaml_entry("caption", &MetaValue::Str(v.clone())),
|
||||
want,
|
||||
"{v:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_quoted_caption_folds_like_a_plain_one() {
|
||||
let caption = "A track: ".to_string() + &"word ".repeat(30) + "end";
|
||||
let entry = yaml_entry("caption", &MetaValue::Str(caption));
|
||||
let lines: Vec<&str> = entry.trim_end().split('\n').collect();
|
||||
assert!(lines[0].starts_with("caption: 'A track: word"));
|
||||
assert!(lines.len() > 1);
|
||||
assert!(lines[1..].iter().all(|l| l.starts_with(" ")));
|
||||
assert!(lines[0].chars().count() > BEST_WIDTH);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn genres_and_unknown_fields_never_reach_the_prompt() {
|
||||
let plan = parse_plan("<think>\nbpm: 90\ngenres: jazz\nmood: calm\n</think>");
|
||||
assert!(plan.contains_key("genres"));
|
||||
assert!(!plan.contains_key("mood"));
|
||||
assert_eq!(format_plan(&plan), "<think>\nbpm: 90\n</think>");
|
||||
}
|
||||
}
|
||||
@@ -14,11 +14,12 @@
|
||||
//!
|
||||
//! Reference: <https://huggingface.co/ACE-Step/ACE-Step-v1-3.5B>
|
||||
|
||||
use super::linear::{Linear, linear_b, linear_no_bias};
|
||||
use candle_core::{D, DType, Device, Module, Result, Tensor};
|
||||
use candle_nn::{
|
||||
Activation, Conv1d, Conv1dConfig, ConvTranspose1d, ConvTranspose1dConfig, VarBuilder,
|
||||
};
|
||||
use candle_transformers::models::with_tracing::{Linear, RmsNorm, linear_b, linear_no_bias};
|
||||
use candle_transformers::models::with_tracing::RmsNorm;
|
||||
use candle_transformers::utils::repeat_kv;
|
||||
|
||||
use super::AceStepConfig;
|
||||
|
||||
@@ -0,0 +1,579 @@
|
||||
//! Constrained decoding of the planner's chain of thought (#384).
|
||||
//!
|
||||
//! A port of upstream's `MetadataConstrainedLogitsProcessor` as the CoT
|
||||
//! phase uses it: genres skipped, caption and language generated, no
|
||||
//! user-provided fields, and generation stopped where `</think>` would
|
||||
//! begin (`stop_at_reasoning`). It walks
|
||||
//!
|
||||
//! ```text
|
||||
//! <think> \n bpm: <30..300> \n caption: <free text> \n duration: <10..600> \n
|
||||
//! keyscale: <note[acc] major|minor> \n language: <code> \n timesignature: <2|3|4|6> \n
|
||||
//! ```
|
||||
//!
|
||||
//! and ends in EOS. Each step it masks the logits to what the current state
|
||||
//! allows; after sampling it advances on the token chosen. Field values are
|
||||
//! matched as token-id sequences against prefix trees built from how the
|
||||
//! tokenizer actually splits `"bpm: 120"` and the like, as upstream does.
|
||||
//!
|
||||
//! Reference: `acestep/constrained_logits_processor.py` at ACE-Step 1.5
|
||||
//! `ca1e85fe` (`_process_single_sequence`, `update_state`, and the prefix
|
||||
//! tree builders).
|
||||
|
||||
use std::collections::{BTreeSet, HashMap};
|
||||
|
||||
use candle_core::Result;
|
||||
|
||||
/// Upstream's `BPM_MIN..=BPM_MAX`.
|
||||
const BPM: (u32, u32) = (30, 300);
|
||||
/// Upstream's `DURATION_MIN..=DURATION_MAX`.
|
||||
const DURATION: (u32, u32) = (10, 600);
|
||||
const TIME_SIGNATURES: [&str; 4] = ["2", "3", "4", "6"];
|
||||
const KEYSCALE_NOTES: [&str; 7] = ["A", "B", "C", "D", "E", "F", "G"];
|
||||
const KEYSCALE_ACCIDENTALS: [&str; 5] = ["", "#", "b", "♯", "♭"];
|
||||
const KEYSCALE_MODES: [&str; 2] = ["major", "minor"];
|
||||
/// Upstream's `VALID_LANGUAGES`, `unknown` included.
|
||||
pub const LANGUAGES: [&str; 51] = [
|
||||
"ar", "az", "bg", "bn", "ca", "cs", "da", "de", "el", "en", "es", "fa", "fi", "fr", "he", "hi",
|
||||
"hr", "ht", "hu", "id", "is", "it", "ja", "ko", "la", "lt", "ms", "ne", "nl", "no", "pa", "pl",
|
||||
"pt", "ro", "ru", "sa", "sk", "sr", "sv", "sw", "ta", "te", "th", "tl", "tr", "uk", "ur", "vi",
|
||||
"yue", "zh", "unknown",
|
||||
];
|
||||
/// Caption tokens before the caption is forced to end.
|
||||
const MAX_CAPTION_TOKENS: usize = 512;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum State {
|
||||
ThinkTag,
|
||||
NewlineAfterThink,
|
||||
BpmName,
|
||||
BpmValue,
|
||||
CaptionName,
|
||||
CaptionValue,
|
||||
DurationName,
|
||||
DurationValue,
|
||||
KeyscaleName,
|
||||
KeyscaleValue,
|
||||
LanguageName,
|
||||
LanguageValue,
|
||||
TimesigName,
|
||||
TimesigValue,
|
||||
ThinkEndTag,
|
||||
Completed,
|
||||
}
|
||||
|
||||
impl State {
|
||||
fn fixed(self) -> Option<&'static str> {
|
||||
Some(match self {
|
||||
Self::ThinkTag => "<think>",
|
||||
Self::NewlineAfterThink => "\n",
|
||||
Self::BpmName => "bpm:",
|
||||
Self::CaptionName => "caption:",
|
||||
Self::DurationName => "duration:",
|
||||
Self::KeyscaleName => "keyscale:",
|
||||
Self::LanguageName => "language:",
|
||||
Self::TimesigName => "timesignature:",
|
||||
Self::ThinkEndTag => "</think>",
|
||||
_ => return None,
|
||||
})
|
||||
}
|
||||
|
||||
/// Field order with genres skipped.
|
||||
fn next(self) -> Self {
|
||||
match self {
|
||||
Self::ThinkTag => Self::NewlineAfterThink,
|
||||
Self::NewlineAfterThink => Self::BpmName,
|
||||
Self::BpmName => Self::BpmValue,
|
||||
Self::BpmValue => Self::CaptionName,
|
||||
Self::CaptionName => Self::CaptionValue,
|
||||
Self::CaptionValue => Self::DurationName,
|
||||
Self::DurationName => Self::DurationValue,
|
||||
Self::DurationValue => Self::KeyscaleName,
|
||||
Self::KeyscaleName => Self::KeyscaleValue,
|
||||
Self::KeyscaleValue => Self::LanguageName,
|
||||
Self::LanguageName => Self::LanguageValue,
|
||||
Self::LanguageValue => Self::TimesigName,
|
||||
Self::TimesigName => Self::TimesigValue,
|
||||
Self::TimesigValue => Self::ThinkEndTag,
|
||||
Self::ThinkEndTag | Self::Completed => Self::Completed,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Token-id prefix → the tokens that may follow it.
|
||||
type PrefixTree = HashMap<Vec<u32>, BTreeSet<u32>>;
|
||||
|
||||
/// The tokenizer, as the state machine needs it.
|
||||
pub trait Vocab {
|
||||
/// Encode without special tokens added (special tokens in the text
|
||||
/// still encode to their ids).
|
||||
fn encode(&self, text: &str) -> Result<Vec<u32>>;
|
||||
/// One token's text, special tokens kept.
|
||||
fn decode(&self, token: u32) -> Result<String>;
|
||||
}
|
||||
|
||||
/// Build a prefix tree for `values` as they tokenize after `"{field}: "`,
|
||||
/// keyed from where `"{field}:"` ends. A value whose tokens do not start
|
||||
/// with the field's own tokens is skipped, as upstream does.
|
||||
fn prefix_tree(
|
||||
vocab: &dyn Vocab,
|
||||
field: &str,
|
||||
values: &[String],
|
||||
newline: u32,
|
||||
first_must_be_note: bool,
|
||||
) -> Result<PrefixTree> {
|
||||
let context = vocab.encode(&format!("{field}:"))?;
|
||||
let mut tree = PrefixTree::new();
|
||||
for value in values {
|
||||
let full = vocab.encode(&format!("{field}: {value}"))?;
|
||||
if full.len() < context.len() || full[..context.len()] != context[..] {
|
||||
continue;
|
||||
}
|
||||
let ids = &full[context.len()..];
|
||||
if ids.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if first_must_be_note {
|
||||
let first = vocab.decode(ids[0])?;
|
||||
let c = first
|
||||
.trim_start()
|
||||
.chars()
|
||||
.next()
|
||||
.map(|c| c.to_ascii_uppercase());
|
||||
if !c.is_some_and(|c| "ABCDEFG".contains(c)) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
for i in 0..=ids.len() {
|
||||
let entry = tree.entry(ids[..i].to_vec()).or_default();
|
||||
entry.insert(if i < ids.len() { ids[i] } else { newline });
|
||||
}
|
||||
}
|
||||
Ok(tree)
|
||||
}
|
||||
|
||||
/// Keep only `allowed`; with nothing allowed, every logit is `-inf`.
|
||||
fn whitelist(logits: &mut [f32], allowed: &[u32]) {
|
||||
let keep: Vec<(usize, f32)> = allowed
|
||||
.iter()
|
||||
.filter_map(|&t| logits.get(t as usize).map(|&v| (t as usize, v)))
|
||||
.collect();
|
||||
logits.fill(f32::NEG_INFINITY);
|
||||
for (i, v) in keep {
|
||||
logits[i] = v;
|
||||
}
|
||||
}
|
||||
|
||||
/// `torch.argmax`: the first maximum.
|
||||
fn argmax(logits: &[f32]) -> u32 {
|
||||
let mut best = 0;
|
||||
for (i, &v) in logits.iter().enumerate() {
|
||||
if v > logits[best] {
|
||||
best = i;
|
||||
}
|
||||
}
|
||||
best as u32
|
||||
}
|
||||
|
||||
/// The CoT phase's constraint, one per generation.
|
||||
pub struct MetaFsm {
|
||||
state: State,
|
||||
position: usize,
|
||||
acc_ids: Vec<u32>,
|
||||
acc_value: String,
|
||||
caption_after_newline: bool,
|
||||
caption_tokens: usize,
|
||||
caption_ending: bool,
|
||||
pending_field: String,
|
||||
|
||||
newline: u32,
|
||||
eos: u32,
|
||||
backtick: Option<u32>,
|
||||
audio_codes: (u32, u32),
|
||||
bpm: PrefixTree,
|
||||
duration: PrefixTree,
|
||||
timesig: PrefixTree,
|
||||
keyscale: PrefixTree,
|
||||
language: PrefixTree,
|
||||
}
|
||||
|
||||
impl MetaFsm {
|
||||
/// `eos` is `<|im_end|>`; `audio_codes` the inclusive id range of the
|
||||
/// `<|audio_code_N|>` tokens, which a caption may never contain.
|
||||
pub fn new(vocab: &dyn Vocab, eos: u32, audio_codes: (u32, u32)) -> Result<Self> {
|
||||
let last = |s: &str| -> Result<Option<u32>> { Ok(vocab.encode(s)?.last().copied()) };
|
||||
let newline =
|
||||
last("\n")?.ok_or_else(|| candle_core::Error::Msg("newline has no token".into()))?;
|
||||
let range = |(lo, hi): (u32, u32)| (lo..=hi).map(|v| v.to_string()).collect::<Vec<_>>();
|
||||
let mut keyscales = Vec::new();
|
||||
for note in KEYSCALE_NOTES {
|
||||
for acc in KEYSCALE_ACCIDENTALS {
|
||||
for mode in KEYSCALE_MODES {
|
||||
keyscales.push(format!("{note}{acc} {mode}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
let strings = |v: &[&str]| v.iter().map(|s| s.to_string()).collect::<Vec<_>>();
|
||||
Ok(Self {
|
||||
state: State::ThinkTag,
|
||||
position: 0,
|
||||
acc_ids: Vec::new(),
|
||||
acc_value: String::new(),
|
||||
caption_after_newline: false,
|
||||
caption_tokens: 0,
|
||||
caption_ending: false,
|
||||
pending_field: String::new(),
|
||||
newline,
|
||||
eos,
|
||||
backtick: last("`")?,
|
||||
audio_codes,
|
||||
bpm: prefix_tree(vocab, "bpm", &range(BPM), newline, false)?,
|
||||
duration: prefix_tree(vocab, "duration", &range(DURATION), newline, false)?,
|
||||
timesig: prefix_tree(
|
||||
vocab,
|
||||
"timesignature",
|
||||
&strings(&TIME_SIGNATURES),
|
||||
newline,
|
||||
false,
|
||||
)?,
|
||||
keyscale: prefix_tree(vocab, "keyscale", &keyscales, newline, true)?,
|
||||
language: prefix_tree(vocab, "language", &strings(&LANGUAGES), newline, false)?,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn is_completed(&self) -> bool {
|
||||
self.state == State::Completed
|
||||
}
|
||||
|
||||
/// The longest prefix of what remains of a fixed string that is one
|
||||
/// token; failing that, the first token of the remainder.
|
||||
fn fixed_token(&self, vocab: &dyn Vocab, fixed: &str) -> Result<Option<u32>> {
|
||||
let rest: Vec<char> = fixed.chars().skip(self.position).collect();
|
||||
if rest.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
for end in (1..=rest.len()).rev() {
|
||||
let ids = vocab.encode(&rest[..end].iter().collect::<String>())?;
|
||||
if ids.len() == 1 {
|
||||
return Ok(Some(ids[0]));
|
||||
}
|
||||
}
|
||||
Ok(vocab
|
||||
.encode(&rest.iter().collect::<String>())?
|
||||
.first()
|
||||
.copied())
|
||||
}
|
||||
|
||||
fn transition(&mut self) {
|
||||
self.state = self.state.next();
|
||||
self.position = 0;
|
||||
self.acc_value.clear();
|
||||
self.acc_ids.clear();
|
||||
self.caption_after_newline = false;
|
||||
self.caption_tokens = 0;
|
||||
self.caption_ending = false;
|
||||
self.pending_field.clear();
|
||||
}
|
||||
|
||||
fn tree(&self) -> &PrefixTree {
|
||||
match self.state {
|
||||
State::BpmValue => &self.bpm,
|
||||
State::DurationValue => &self.duration,
|
||||
State::TimesigValue => &self.timesig,
|
||||
State::KeyscaleValue => &self.keyscale,
|
||||
_ => &self.language,
|
||||
}
|
||||
}
|
||||
|
||||
fn allows_newline(&self) -> bool {
|
||||
self.tree()
|
||||
.get(&self.acc_ids)
|
||||
.is_some_and(|s| s.contains(&self.newline))
|
||||
}
|
||||
|
||||
fn continuations(&self) -> Vec<u32> {
|
||||
self.tree()
|
||||
.get(&self.acc_ids)
|
||||
.map(|s| s.iter().copied().collect())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Mask `logits` to what the current state allows.
|
||||
pub fn constrain(&mut self, vocab: &dyn Vocab, logits: &mut [f32]) -> Result<()> {
|
||||
if self.state == State::Completed {
|
||||
return Ok(());
|
||||
}
|
||||
if let Some(fixed) = self.state.fixed() {
|
||||
if self.state == State::ThinkEndTag {
|
||||
// stop_at_reasoning: end where `</think>` would begin.
|
||||
whitelist(logits, &[self.eos]);
|
||||
return Ok(());
|
||||
}
|
||||
match self.fixed_token(vocab, fixed)? {
|
||||
Some(t) => whitelist(logits, &[t]),
|
||||
None => {
|
||||
self.transition();
|
||||
if self.state.fixed().is_none() {
|
||||
logits.fill(0.0);
|
||||
return self.constrain(vocab, logits);
|
||||
}
|
||||
}
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
match self.state {
|
||||
State::BpmValue | State::DurationValue => {
|
||||
let mut allowed = self.continuations();
|
||||
if self.allows_newline() {
|
||||
allowed.push(self.newline);
|
||||
}
|
||||
whitelist(logits, &allowed);
|
||||
}
|
||||
State::CaptionValue => {
|
||||
if self.caption_after_newline {
|
||||
self.caption_after_newline = false;
|
||||
let top = vocab.decode(argmax(logits))?;
|
||||
if top.chars().next().is_some_and(|c| c != ' ' && c != '\t') {
|
||||
// The model is starting the next field's name: let
|
||||
// it, and pick the field up from what it writes.
|
||||
self.caption_ending = true;
|
||||
self.pending_field.clear();
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
if self.caption_ending {
|
||||
return Ok(());
|
||||
}
|
||||
if let Some(b) = self.backtick
|
||||
&& let Some(v) = logits.get_mut(b as usize)
|
||||
{
|
||||
*v = f32::NEG_INFINITY;
|
||||
}
|
||||
let (lo, hi) = self.audio_codes;
|
||||
for v in logits.iter_mut().take(hi as usize + 1).skip(lo as usize) {
|
||||
*v = f32::NEG_INFINITY;
|
||||
}
|
||||
if self.caption_tokens >= MAX_CAPTION_TOKENS {
|
||||
whitelist(logits, &[self.newline]);
|
||||
}
|
||||
}
|
||||
State::KeyscaleValue | State::TimesigValue => {
|
||||
if self.allows_newline() {
|
||||
whitelist(logits, &[self.newline]);
|
||||
} else {
|
||||
let allowed = self.continuations();
|
||||
if allowed.is_empty() && self.state == State::KeyscaleValue {
|
||||
whitelist(logits, &[self.newline]);
|
||||
} else {
|
||||
whitelist(logits, &allowed);
|
||||
}
|
||||
}
|
||||
}
|
||||
State::LanguageValue => {
|
||||
if self.acc_ids.is_empty() {
|
||||
// The single best first token among every language's.
|
||||
let candidates = self.continuations();
|
||||
match candidates.iter().copied().reduce(|best, t| {
|
||||
let v = |t: u32| logits.get(t as usize).copied().unwrap_or(f32::MIN);
|
||||
if v(t) > v(best) { t } else { best }
|
||||
}) {
|
||||
Some(t) => whitelist(logits, &[t]),
|
||||
None => whitelist(logits, &[self.newline]),
|
||||
}
|
||||
} else if self.allows_newline() {
|
||||
whitelist(logits, &[self.newline]);
|
||||
} else {
|
||||
let allowed = self.continuations();
|
||||
if allowed.is_empty() {
|
||||
whitelist(logits, &[self.newline]);
|
||||
} else {
|
||||
whitelist(logits, &allowed);
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Advance on the token that was chosen.
|
||||
pub fn update(&mut self, vocab: &dyn Vocab, token: u32) -> Result<()> {
|
||||
if self.state == State::Completed {
|
||||
return Ok(());
|
||||
}
|
||||
let text = vocab.decode(token)?;
|
||||
if let Some(fixed) = self.state.fixed() {
|
||||
self.position += text.chars().count();
|
||||
if self.state == State::ThinkEndTag && token == self.eos {
|
||||
self.state = State::Completed;
|
||||
} else if self.position >= fixed.chars().count() {
|
||||
self.transition();
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
match self.state {
|
||||
State::BpmValue
|
||||
| State::DurationValue
|
||||
| State::TimesigValue
|
||||
| State::KeyscaleValue
|
||||
| State::LanguageValue => {
|
||||
if token == self.newline {
|
||||
self.transition();
|
||||
} else {
|
||||
self.acc_ids.push(token);
|
||||
self.acc_value.push_str(&text);
|
||||
}
|
||||
}
|
||||
State::CaptionValue => {
|
||||
self.caption_tokens += 1;
|
||||
self.acc_value.push_str(&text);
|
||||
self.caption_after_newline = text.contains('\n');
|
||||
if self.caption_ending {
|
||||
self.pending_field.push_str(&text);
|
||||
if text.contains(':') {
|
||||
let name = self
|
||||
.pending_field
|
||||
.trim()
|
||||
.trim_end_matches(':')
|
||||
.trim()
|
||||
.to_lowercase();
|
||||
let target = match name.as_str() {
|
||||
"duration" => Some(State::DurationValue),
|
||||
"keyscale" => Some(State::KeyscaleValue),
|
||||
"language" => Some(State::LanguageValue),
|
||||
"timesignature" => Some(State::TimesigValue),
|
||||
_ => None,
|
||||
};
|
||||
match target {
|
||||
Some(s) => {
|
||||
self.transition();
|
||||
self.state = s;
|
||||
}
|
||||
None => self.transition(),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// A vocabulary that splits text into characters, except that a space
|
||||
/// joins the character after it (as `" G"` is one token in the real
|
||||
/// one), with `<think>`, `</think>` and `<|im_end|>` as specials.
|
||||
/// Enough to drive every state; the real tokenizer is exercised by the
|
||||
/// planner parity test.
|
||||
struct Chars;
|
||||
const SPECIALS: [&str; 3] = ["<think>", "</think>", "<|im_end|>"];
|
||||
const BASE: u32 = 1000;
|
||||
const SPACED: u32 = 3000;
|
||||
|
||||
impl Vocab for Chars {
|
||||
fn encode(&self, text: &str) -> Result<Vec<u32>> {
|
||||
if let Some(i) = SPECIALS.iter().position(|s| *s == text) {
|
||||
return Ok(vec![BASE + i as u32]);
|
||||
}
|
||||
let chars: Vec<char> = text.chars().collect();
|
||||
let mut ids = Vec::new();
|
||||
let mut i = 0;
|
||||
while i < chars.len() {
|
||||
if chars[i] == ' '
|
||||
&& i + 1 < chars.len()
|
||||
&& (chars[i + 1] as u32) < 128
|
||||
&& chars[i + 1] != ' '
|
||||
{
|
||||
ids.push(SPACED + chars[i + 1] as u32);
|
||||
i += 2;
|
||||
} else {
|
||||
ids.push(chars[i] as u32);
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
Ok(ids)
|
||||
}
|
||||
fn decode(&self, token: u32) -> Result<String> {
|
||||
Ok(match token {
|
||||
t if (BASE..BASE + 3).contains(&t) => SPECIALS[(t - BASE) as usize].to_string(),
|
||||
t if (SPACED..SPACED + 128).contains(&t) => {
|
||||
format!(" {}", char::from_u32(t - SPACED).unwrap_or('?'))
|
||||
}
|
||||
t => char::from_u32(t).unwrap_or('?').to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Drive the machine with a model that always prefers the next
|
||||
/// character of `wanted`; return what it generates.
|
||||
fn run(wanted: &str) -> String {
|
||||
let v = Chars;
|
||||
let eos = BASE + 2;
|
||||
let mut fsm = MetaFsm::new(&v, eos, (2000, 2010)).unwrap();
|
||||
let want: Vec<u32> = {
|
||||
let mut ids = vec![BASE];
|
||||
ids.extend(v.encode(wanted).unwrap());
|
||||
ids
|
||||
};
|
||||
let mut out = String::new();
|
||||
for step in 0..400 {
|
||||
let mut logits = vec![0f32; 4000];
|
||||
if let Some(&t) = want.get(step) {
|
||||
logits[t as usize] = 10.0;
|
||||
}
|
||||
fsm.constrain(&v, &mut logits).unwrap();
|
||||
let t = argmax(&logits);
|
||||
fsm.update(&v, t).unwrap();
|
||||
if t == eos {
|
||||
out.push_str("<EOS>");
|
||||
break;
|
||||
}
|
||||
out.push_str(&v.decode(t).unwrap());
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_well_formed_plan_passes_through_and_stops_before_the_close_tag() {
|
||||
let plan = "\nbpm: 120\ncaption: A calm piano piece\n with soft strings\nduration: 30\nkeyscale: C major\nlanguage: en\ntimesignature: 4\n";
|
||||
assert_eq!(run(plan), format!("<think>{plan}<EOS>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn out_of_range_and_off_vocabulary_values_are_steered() {
|
||||
// bpm 999 is out of range; keyscale H is not a note; language xx
|
||||
// is not a code. The model's preference is followed only where the
|
||||
// grammar allows it.
|
||||
let plan = "\nbpm: 999\ncaption: x\nduration: 30\nkeyscale: H major\nlanguage: xx\ntimesignature: 5\n";
|
||||
let got = run(plan);
|
||||
assert!(got.starts_with("<think>\nbpm: 99"), "{got}");
|
||||
let bpm: u32 = got["<think>\nbpm: ".len()..]
|
||||
.split('\n')
|
||||
.next()
|
||||
.unwrap()
|
||||
.parse()
|
||||
.unwrap();
|
||||
assert!((BPM.0..=BPM.1).contains(&bpm), "{got}");
|
||||
assert!(got.contains("\nkeyscale: A"), "{got}");
|
||||
let lang = got
|
||||
.split("language: ")
|
||||
.nth(1)
|
||||
.unwrap()
|
||||
.split('\n')
|
||||
.next()
|
||||
.unwrap();
|
||||
assert!(LANGUAGES.contains(&lang), "{got}");
|
||||
let ts = got
|
||||
.split("timesignature: ")
|
||||
.nth(1)
|
||||
.unwrap()
|
||||
.split('\n')
|
||||
.next()
|
||||
.unwrap();
|
||||
assert!(TIME_SIGNATURES.contains(&ts), "{got}");
|
||||
assert!(got.ends_with("<EOS>"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,162 @@
|
||||
//! A linear layer that is either dense or quantized in situ (#385).
|
||||
//!
|
||||
//! The vendored modules build their projections through `linear_b` /
|
||||
//! `linear_no_bias`, the same names `candle_transformers::with_tracing`
|
||||
//! exports, so switching them over is an import change and the vendored
|
||||
//! code stays diffable against upstream.
|
||||
//!
|
||||
//! Whether a layer is quantized is decided while it is built, by
|
||||
//! [`with_quant`]. A scope rather than a parameter, because the setting
|
||||
//! would otherwise have to be threaded through every constructor of the
|
||||
//! DiT, the condition encoder, the text encoder and the planner. Loads
|
||||
//! run on one thread (the device worker), so the scope cannot leak into
|
||||
//! another model's load.
|
||||
//!
|
||||
//! A quantized layer runs through candle's `QMatMul`, which wants f32
|
||||
//! activations: build the component with an f32 `VarBuilder`.
|
||||
|
||||
use candle_core::quantized::{GgmlDType, QMatMul};
|
||||
use candle_core::{Module, Result, Tensor};
|
||||
use candle_nn::VarBuilder;
|
||||
use std::cell::Cell;
|
||||
use std::sync::Arc;
|
||||
|
||||
thread_local! {
|
||||
static QUANT: Cell<Option<GgmlDType>> = const { Cell::new(None) };
|
||||
}
|
||||
|
||||
/// Build with `quant` applied to every linear layer `f` creates. A layer
|
||||
/// whose input width does not divide the block size stays dense.
|
||||
pub fn with_quant<R>(quant: Option<GgmlDType>, f: impl FnOnce() -> R) -> R {
|
||||
let previous = QUANT.with(|q| q.replace(quant));
|
||||
let out = f();
|
||||
QUANT.with(|q| q.set(previous));
|
||||
out
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum Linear {
|
||||
Dense(candle_nn::Linear),
|
||||
Quantized {
|
||||
weight: QMatMul,
|
||||
bias: Option<Tensor>,
|
||||
},
|
||||
}
|
||||
|
||||
impl Linear {
|
||||
/// A dense layer from existing weights (tied heads, identity
|
||||
/// projections). Never quantized: the weight is shared or synthetic.
|
||||
pub fn from_weights(weight: Tensor, bias: Option<Tensor>) -> Self {
|
||||
Self::Dense(candle_nn::Linear::new(weight, bias))
|
||||
}
|
||||
|
||||
fn build(weight: Tensor, bias: Option<Tensor>) -> Result<Self> {
|
||||
let quant = QUANT.with(Cell::get);
|
||||
let width = weight.dims().last().copied().unwrap_or(0);
|
||||
match quant {
|
||||
Some(dtype) if width > 0 && width.is_multiple_of(dtype.block_size()) => {
|
||||
let qt = crate::harness::tp::isq::quantize_parallel(&weight, dtype)
|
||||
.map_err(|e| candle_core::Error::Msg(format!("{e:#}")))?;
|
||||
Ok(Self::Quantized {
|
||||
weight: QMatMul::from_arc(Arc::new(qt))?,
|
||||
bias,
|
||||
})
|
||||
}
|
||||
_ => Ok(Self::Dense(candle_nn::Linear::new(weight, bias))),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_quantized(&self) -> bool {
|
||||
matches!(self, Self::Quantized { .. })
|
||||
}
|
||||
}
|
||||
|
||||
impl Module for Linear {
|
||||
fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
match self {
|
||||
Self::Dense(l) => l.forward(x),
|
||||
Self::Quantized { weight, bias } => {
|
||||
// The quantized kernels read a contiguous (batch, rows, k)
|
||||
// or (rows, k) input.
|
||||
let y = match x.dims() {
|
||||
[a, b, c, k] => weight
|
||||
.forward(&x.contiguous()?.reshape((a * b, *c, *k))?)?
|
||||
.reshape((*a, *b, *c, ()))?,
|
||||
_ => weight.forward(&x.contiguous()?)?,
|
||||
};
|
||||
match bias {
|
||||
Some(b) => y.broadcast_add(b),
|
||||
None => Ok(y),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn linear_b(d1: usize, d2: usize, bias: bool, vb: VarBuilder) -> Result<Linear> {
|
||||
let weight = vb.get((d2, d1), "weight")?;
|
||||
let bias = if bias {
|
||||
Some(vb.get(d2, "bias")?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
Linear::build(weight, bias)
|
||||
}
|
||||
|
||||
pub fn linear_no_bias(d1: usize, d2: usize, vb: VarBuilder) -> Result<Linear> {
|
||||
linear_b(d1, d2, false, vb)
|
||||
}
|
||||
|
||||
pub fn linear(d1: usize, d2: usize, vb: VarBuilder) -> Result<Linear> {
|
||||
linear_b(d1, d2, true, vb)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use candle_core::{DType, Device};
|
||||
use std::collections::HashMap;
|
||||
|
||||
fn vb_with(weight: Tensor) -> VarBuilder<'static> {
|
||||
let mut map = HashMap::new();
|
||||
map.insert("weight".to_string(), weight);
|
||||
VarBuilder::from_tensors(map, DType::F32, &Device::Cpu)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quantization_is_scoped_and_skips_widths_that_do_not_tile() -> Result<()> {
|
||||
let w = Tensor::randn(0f32, 1.0, (64, 64), &Device::Cpu)?;
|
||||
let dense = linear_no_bias(64, 64, vb_with(w.clone()))?;
|
||||
assert!(!dense.is_quantized(), "no scope, no quantization");
|
||||
|
||||
let q = with_quant(Some(GgmlDType::Q8_0), || {
|
||||
linear_no_bias(64, 64, vb_with(w.clone()))
|
||||
})?;
|
||||
assert!(q.is_quantized());
|
||||
|
||||
let narrow = Tensor::randn(0f32, 1.0, (64, 6), &Device::Cpu)?;
|
||||
let kept = with_quant(Some(GgmlDType::Q8_0), || {
|
||||
linear_no_bias(6, 64, vb_with(narrow))
|
||||
})?;
|
||||
assert!(!kept.is_quantized(), "6 does not tile a 32-block");
|
||||
|
||||
let after = linear_no_bias(64, 64, vb_with(w))?;
|
||||
assert!(!after.is_quantized(), "the scope must end with the call");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn q8_output_tracks_the_dense_layer() -> Result<()> {
|
||||
let w = Tensor::randn(0f32, 0.05, (128, 256), &Device::Cpu)?;
|
||||
let x = Tensor::randn(0f32, 1.0, (3, 256), &Device::Cpu)?;
|
||||
let dense = linear_no_bias(256, 128, vb_with(w.clone()))?.forward(&x)?;
|
||||
let q = with_quant(Some(GgmlDType::Q8_0), || {
|
||||
linear_no_bias(256, 128, vb_with(w))
|
||||
})?
|
||||
.forward(&x)?;
|
||||
let err = (dense.clone() - q)?.abs()?.max_all()?.to_scalar::<f32>()?;
|
||||
let scale = dense.abs()?.max_all()?.to_scalar::<f32>()?;
|
||||
assert!(err / scale < 0.02, "q8_0 relative error {}", err / scale);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -30,7 +30,10 @@
|
||||
//! - [GitHub Repository](https://github.com/ACE-Step/ACE-Step-1.5)
|
||||
|
||||
pub mod condition;
|
||||
pub mod cot;
|
||||
pub mod dit;
|
||||
pub mod fsm;
|
||||
pub mod linear;
|
||||
pub mod lm;
|
||||
pub mod model;
|
||||
pub mod pipeline;
|
||||
|
||||
@@ -18,6 +18,17 @@ use anyhow::Result;
|
||||
use candle_core::{DType, Device, Tensor};
|
||||
use candle_nn::VarBuilder;
|
||||
|
||||
/// The DiT's task instruction for a text-to-music render (upstream's
|
||||
/// `TASK_INSTRUCTIONS["text2music"]`).
|
||||
pub const TEXT2MUSIC_INSTRUCTION: &str =
|
||||
"Fill the audio semantic mask based on the given conditions:";
|
||||
/// The instruction once audio codes condition the render. Upstream turns a
|
||||
/// text-to-music request into a cover when the planner supplies codes
|
||||
/// (`_resolve_generate_music_task`), and so does a reference cover.
|
||||
pub const COVER_INSTRUCTION: &str = "Generate audio semantic tokens based on the given conditions:";
|
||||
/// Upstream's `TASK_INSTRUCTIONS["repaint"]`.
|
||||
pub const REPAINT_INSTRUCTION: &str = "Repaint the mask area based on the given conditions:";
|
||||
|
||||
/// Metadata from LM CoT output, used to enrich the DiT text prompt.
|
||||
///
|
||||
/// When the LM generates chain-of-thought metadata (BPM, keyscale, etc.),
|
||||
@@ -94,6 +105,21 @@ pub struct AudioOutput {
|
||||
pub sample_rate: u32,
|
||||
/// The denoised latents `(1, T, acoustic_dim)` the VAE decoded.
|
||||
pub latents: Tensor,
|
||||
/// Wall time per phase. Each phase ends with a device synchronize, so
|
||||
/// queued GPU work is charged to the phase that queued it.
|
||||
pub timing: PhaseTiming,
|
||||
}
|
||||
|
||||
/// Wall time of one render's phases, in milliseconds.
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
pub struct PhaseTiming {
|
||||
pub encode_ms: u64,
|
||||
pub denoise_ms: u64,
|
||||
pub decode_ms: u64,
|
||||
}
|
||||
|
||||
fn elapsed_ms(since: std::time::Instant) -> u64 {
|
||||
since.elapsed().as_millis() as u64
|
||||
}
|
||||
|
||||
/// High-level pipeline: text prompt → stereo audio.
|
||||
@@ -107,8 +133,16 @@ pub struct AceStepPipeline {
|
||||
silence_latent: Tensor,
|
||||
dit_config: AceStepConfig,
|
||||
vae_config: VaeConfig,
|
||||
/// The DiT's device and dtype. Conditioning, latents and the planner's
|
||||
/// hints live here.
|
||||
device: Device,
|
||||
dtype: DType,
|
||||
/// Where the text encoder and the VAE run, which may differ from the
|
||||
/// DiT's (#385). Only their small outputs cross between devices: the
|
||||
/// prompt's hidden states and the denoised latents.
|
||||
text_device: Device,
|
||||
vae_device: Device,
|
||||
vae_dtype: DType,
|
||||
}
|
||||
|
||||
struct LmState {
|
||||
@@ -119,7 +153,8 @@ struct LmState {
|
||||
impl AceStepPipeline {
|
||||
/// Load all components.
|
||||
///
|
||||
/// - `dit_vb` / `vae_vb` / `text_vb`: VarBuilders for each model
|
||||
/// - `dit_vb` / `vae_vb` / `text_vb`: VarBuilders for each model. Each
|
||||
/// component runs on its builder's device and dtype.
|
||||
/// - `text_config`: Qwen3 config (for the embedding model, NOT the LM)
|
||||
/// - `tokenizer_path`: `tokenizer.json` for the text encoder
|
||||
/// - `silence_latent`: `(1, T, acoustic_dim)` pre-loaded
|
||||
@@ -138,6 +173,9 @@ impl AceStepPipeline {
|
||||
) -> Result<Self> {
|
||||
let device = dit_vb.device().clone();
|
||||
let dtype = dit_vb.dtype();
|
||||
let text_device = text_vb.device().clone();
|
||||
let vae_device = vae_vb.device().clone();
|
||||
let vae_dtype = vae_vb.dtype();
|
||||
|
||||
let text_encoder = qwen3::Model::new(text_config, text_vb)?;
|
||||
let tokenizer = tokenizers::Tokenizer::from_file(tokenizer_path)
|
||||
@@ -162,6 +200,9 @@ impl AceStepPipeline {
|
||||
vae_config: vae_config.clone(),
|
||||
device,
|
||||
dtype,
|
||||
text_device,
|
||||
vae_device,
|
||||
vae_dtype,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -178,18 +219,27 @@ impl AceStepPipeline {
|
||||
lyrics: &str,
|
||||
params: &GenerationParams,
|
||||
) -> Result<AudioOutput> {
|
||||
let t = std::time::Instant::now();
|
||||
let instruction = if params.lm_hints_25hz.is_some() {
|
||||
COVER_INSTRUCTION
|
||||
} else {
|
||||
TEXT2MUSIC_INSTRUCTION
|
||||
};
|
||||
let (text_hs, text_mask, lyric_hs, lyric_mask) = self.encode_text(
|
||||
instruction,
|
||||
prompt,
|
||||
lyrics,
|
||||
params.duration_secs,
|
||||
params.lm_metadata.as_ref(),
|
||||
params.language.as_deref(),
|
||||
)?;
|
||||
self.device.synchronize()?;
|
||||
let encode_ms = elapsed_ms(t);
|
||||
|
||||
let latent_frames = self.latent_frames(params.duration_secs);
|
||||
let timbre_frames = 750.min(self.silence_latent.dim(1)?);
|
||||
|
||||
self.generate_inner(
|
||||
let mut out = self.generate_inner(
|
||||
&text_hs,
|
||||
&text_mask,
|
||||
&lyric_hs,
|
||||
@@ -198,7 +248,9 @@ impl AceStepPipeline {
|
||||
&self.silence_latent.narrow(1, 0, timbre_frames)?,
|
||||
params,
|
||||
latent_frames,
|
||||
)
|
||||
)?;
|
||||
out.timing.encode_ms = encode_ms;
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Generate a cover conditioned on reference audio `(1, 2, T)` at 48kHz.
|
||||
@@ -221,6 +273,7 @@ impl AceStepPipeline {
|
||||
params: &GenerationParams,
|
||||
) -> Result<AudioOutput> {
|
||||
let (text_hs, text_mask, lyric_hs, lyric_mask) = self.encode_text(
|
||||
COVER_INSTRUCTION,
|
||||
prompt,
|
||||
lyrics,
|
||||
params.duration_secs,
|
||||
@@ -232,7 +285,11 @@ impl AceStepPipeline {
|
||||
.vae_encoder
|
||||
.as_ref()
|
||||
.ok_or_else(|| anyhow::anyhow!("VAE encoder not loaded (pass load_encoder=true)"))?;
|
||||
let ref_latent = vae.encode_mean(reference_audio)?.transpose(1, 2)?;
|
||||
let ref_latent = vae
|
||||
.encode_mean(reference_audio)?
|
||||
.transpose(1, 2)?
|
||||
.to_device(&self.device)?
|
||||
.to_dtype(self.dtype)?;
|
||||
|
||||
let latent_frames = self.latent_frames(params.duration_secs);
|
||||
let ref_len = ref_latent.dim(1)?;
|
||||
@@ -246,7 +303,11 @@ impl AceStepPipeline {
|
||||
|
||||
// Timbre: use separate timbre audio if provided, otherwise reference audio
|
||||
let refer_audio_packed = if let Some(timbre) = timbre_audio {
|
||||
let timbre_latent = vae.encode_mean(timbre)?.transpose(1, 2)?;
|
||||
let timbre_latent = vae
|
||||
.encode_mean(timbre)?
|
||||
.transpose(1, 2)?
|
||||
.to_device(&self.device)?
|
||||
.to_dtype(self.dtype)?;
|
||||
let timbre_len = timbre_latent.dim(1)?;
|
||||
let timbre_frames = 750.min(timbre_len);
|
||||
timbre_latent.narrow(1, 0, timbre_frames)?
|
||||
@@ -289,6 +350,7 @@ impl AceStepPipeline {
|
||||
params: &GenerationParams,
|
||||
) -> Result<AudioOutput> {
|
||||
let (text_hs, text_mask, lyric_hs, lyric_mask) = self.encode_text(
|
||||
REPAINT_INSTRUCTION,
|
||||
prompt,
|
||||
lyrics,
|
||||
params.duration_secs,
|
||||
@@ -300,7 +362,11 @@ impl AceStepPipeline {
|
||||
.vae_encoder
|
||||
.as_ref()
|
||||
.ok_or_else(|| anyhow::anyhow!("VAE encoder not loaded (pass load_encoder=true)"))?;
|
||||
let src_latent = vae.encode_mean(source_audio)?.transpose(1, 2)?; // (1, T_lat, D)
|
||||
let src_latent = vae
|
||||
.encode_mean(source_audio)?
|
||||
.transpose(1, 2)?
|
||||
.to_device(&self.device)?
|
||||
.to_dtype(self.dtype)?; // (1, T_lat, D)
|
||||
|
||||
let latent_frames = self.latent_frames(params.duration_secs);
|
||||
let hop = self.vae_config.hop_length() as f64;
|
||||
@@ -394,6 +460,9 @@ impl AceStepPipeline {
|
||||
lm_vb: VarBuilder,
|
||||
lm_tokenizer_path: &std::path::Path,
|
||||
) -> Result<()> {
|
||||
// The planner runs on its builder's device; only its integer codes
|
||||
// reach the DiT.
|
||||
let lm_device = lm_vb.device().clone();
|
||||
let lm_model = qwen3::ModelForCausalLM::new(lm_config, lm_vb)?;
|
||||
let lm_tokenizer = tokenizers::Tokenizer::from_file(lm_tokenizer_path)
|
||||
.map_err(|e| anyhow::anyhow!("LM tokenizer: {e}"))?;
|
||||
@@ -403,7 +472,7 @@ impl AceStepPipeline {
|
||||
.map_err(|e| candle_core::Error::Msg(format!("encode {s:?}: {e}")))?;
|
||||
Ok(enc.get_ids().to_vec())
|
||||
})?;
|
||||
let pipeline = LmPipeline::new(lm_model, token_ids, self.device.clone());
|
||||
let pipeline = LmPipeline::new(lm_model, token_ids, lm_device);
|
||||
self.lm = Some(LmState {
|
||||
pipeline,
|
||||
tokenizer: lm_tokenizer,
|
||||
@@ -427,15 +496,9 @@ impl AceStepPipeline {
|
||||
.as_mut()
|
||||
.ok_or_else(|| anyhow::anyhow!("LM not loaded — call load_lm() first"))?;
|
||||
|
||||
let prompt_text = LmPipeline::format_prompt(prompt, lyrics);
|
||||
let prompt_enc = lm
|
||||
.tokenizer
|
||||
.encode(prompt_text.as_str(), false)
|
||||
.map_err(|e| anyhow::anyhow!("LM encode: {e}"))?;
|
||||
|
||||
let decode_fn = |t: u32| -> candle_core::Result<String> {
|
||||
let decode_fn = |ids: &[u32]| -> candle_core::Result<String> {
|
||||
lm.tokenizer
|
||||
.decode(&[t], false)
|
||||
.decode(ids, false)
|
||||
.map_err(|e| candle_core::Error::Msg(format!("{e}")))
|
||||
};
|
||||
let encode_fn = |s: &str| -> candle_core::Result<Vec<u32>> {
|
||||
@@ -447,7 +510,8 @@ impl AceStepPipeline {
|
||||
};
|
||||
|
||||
let output = lm.pipeline.generate(
|
||||
prompt_enc.get_ids(),
|
||||
prompt,
|
||||
lyrics,
|
||||
duration_secs,
|
||||
lm_config,
|
||||
&decode_fn,
|
||||
@@ -479,6 +543,7 @@ impl AceStepPipeline {
|
||||
/// lyrics under their language. `language` defaults to `"unknown"`,
|
||||
/// upstream's own default when the caller names none.
|
||||
pub fn format_prompts(
|
||||
instruction: &str,
|
||||
prompt: &str,
|
||||
lyrics: &str,
|
||||
duration_secs: f64,
|
||||
@@ -496,7 +561,7 @@ impl AceStepPipeline {
|
||||
.and_then(|m| m.keyscale.as_deref())
|
||||
.unwrap_or("N/A");
|
||||
let caption_text = format!(
|
||||
"# Instruction\nFill the audio semantic mask based on the given conditions:\n\n\
|
||||
"# Instruction\n{instruction}\n\n\
|
||||
# Caption\n{caption}\n\n# Metas\n\
|
||||
- bpm: {bpm}\n- timesignature: {timesig}\n- keyscale: {keyscale}\n- duration: {} seconds\n\
|
||||
<|endoftext|>\n",
|
||||
@@ -534,28 +599,44 @@ impl AceStepPipeline {
|
||||
|
||||
fn encode_text(
|
||||
&mut self,
|
||||
instruction: &str,
|
||||
prompt: &str,
|
||||
lyrics: &str,
|
||||
duration_secs: f64,
|
||||
metadata: Option<&LmMetadata>,
|
||||
language: Option<&str>,
|
||||
) -> Result<(Tensor, Tensor, Tensor, Tensor)> {
|
||||
let (caption_text, lyric_text) =
|
||||
Self::format_prompts(prompt, lyrics, duration_secs, metadata, language);
|
||||
let (caption_text, lyric_text) = Self::format_prompts(
|
||||
instruction,
|
||||
prompt,
|
||||
lyrics,
|
||||
duration_secs,
|
||||
metadata,
|
||||
language,
|
||||
);
|
||||
let encoding = self
|
||||
.tokenizer
|
||||
.encode(caption_text.as_str(), true)
|
||||
.map_err(|e| anyhow::anyhow!("tokenization: {e}"))?;
|
||||
let prompt_tokens = Tensor::new(encoding.get_ids(), &self.device)?.unsqueeze(0)?;
|
||||
let text_hidden_states = self.text_encoder.forward(&prompt_tokens, 0)?;
|
||||
let prompt_tokens = Tensor::new(encoding.get_ids(), &self.text_device)?.unsqueeze(0)?;
|
||||
let text_hidden_states = self
|
||||
.text_encoder
|
||||
.forward(&prompt_tokens, 0)?
|
||||
.to_device(&self.device)?
|
||||
.to_dtype(self.dtype)?;
|
||||
self.text_encoder.clear_kv_cache();
|
||||
|
||||
let lyric_encoding = self
|
||||
.tokenizer
|
||||
.encode(lyric_text.as_str(), true)
|
||||
.map_err(|e| anyhow::anyhow!("lyric tokenization: {e}"))?;
|
||||
let lyric_tokens = Tensor::new(lyric_encoding.get_ids(), &self.device)?.unsqueeze(0)?;
|
||||
let lyric_hidden_states = self.text_encoder.embed_tokens(&lyric_tokens)?;
|
||||
let lyric_tokens =
|
||||
Tensor::new(lyric_encoding.get_ids(), &self.text_device)?.unsqueeze(0)?;
|
||||
let lyric_hidden_states = self
|
||||
.text_encoder
|
||||
.embed_tokens(&lyric_tokens)?
|
||||
.to_device(&self.device)?
|
||||
.to_dtype(self.dtype)?;
|
||||
|
||||
let text_mask = Tensor::new(encoding.get_attention_mask(), &self.device)?
|
||||
.unsqueeze(0)?
|
||||
@@ -638,6 +719,7 @@ impl AceStepPipeline {
|
||||
.num_steps
|
||||
.unwrap_or(if params.is_turbo { 8 } else { 50 });
|
||||
|
||||
let t = std::time::Instant::now();
|
||||
let xt = self.model.generate_latents(
|
||||
&condition,
|
||||
None,
|
||||
@@ -657,17 +739,32 @@ impl AceStepPipeline {
|
||||
self.dtype,
|
||||
)?;
|
||||
|
||||
self.device.synchronize()?;
|
||||
let denoise_ms = elapsed_ms(t);
|
||||
|
||||
let t = std::time::Instant::now();
|
||||
let latents = xt
|
||||
.transpose(1, 2)?
|
||||
.to_device(&self.vae_device)?
|
||||
.to_dtype(self.vae_dtype)?;
|
||||
let audio = self.vae_decoder.tiled_decode(
|
||||
&xt.transpose(1, 2)?,
|
||||
&latents,
|
||||
params.vae_chunk_frames,
|
||||
params.vae_chunk_overlap,
|
||||
)?;
|
||||
let audio = audio.squeeze(0)?;
|
||||
self.vae_device.synchronize()?;
|
||||
let decode_ms = elapsed_ms(t);
|
||||
|
||||
Ok(AudioOutput {
|
||||
audio,
|
||||
sample_rate: self.vae_config.sampling_rate as u32,
|
||||
latents: xt,
|
||||
timing: PhaseTiming {
|
||||
encode_ms: 0,
|
||||
denoise_ms,
|
||||
decode_ms,
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,9 +1,7 @@
|
||||
use super::linear::{Linear, linear_b, linear_no_bias};
|
||||
use candle_core::{DType, Device, Module, Result, Tensor};
|
||||
use candle_nn::{Activation, VarBuilder, kv_cache::ConcatKvCache};
|
||||
use candle_transformers::{
|
||||
models::with_tracing::{Linear, RmsNorm, linear_b, linear_no_bias},
|
||||
utils::repeat_kv,
|
||||
};
|
||||
use candle_transformers::{models::with_tracing::RmsNorm, utils::repeat_kv};
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
|
||||
|
||||
@@ -12,9 +12,10 @@
|
||||
//! - `AttentionPooler` / `AudioTokenizer`: structurally correct, minor
|
||||
//! float precision differences propagate through FSQ rounding.
|
||||
|
||||
use super::linear::{Linear, linear, linear_no_bias};
|
||||
use candle_core::{D, DType, Module, Result, Tensor};
|
||||
use candle_nn::{Activation, VarBuilder};
|
||||
use candle_transformers::models::with_tracing::{Linear, RmsNorm, linear, linear_no_bias};
|
||||
use candle_transformers::models::with_tracing::RmsNorm;
|
||||
use candle_transformers::utils::repeat_kv;
|
||||
use std::sync::Arc;
|
||||
|
||||
@@ -255,23 +256,18 @@ impl AttentionPooler {
|
||||
|
||||
/// Residual Finite Scalar Quantization.
|
||||
///
|
||||
/// Matches `vector_quantize_pytorch.ResidualFSQ` with a single top-level
|
||||
/// `project_in` (dim → codebook_dim) and `project_out` (codebook_dim → dim).
|
||||
/// Each FSQ layer independently quantizes each codebook dimension to its
|
||||
/// corresponding number of levels using `tanh` bounding and rounding.
|
||||
/// Matches `vector_quantize_pytorch.ResidualFSQ` 1.27 as ACE-Step builds it
|
||||
/// (one quantizer, the library defaults): a top-level `project_in`
|
||||
/// (dim → codebook_dim) and `project_out` (codebook_dim → dim), a soft clamp
|
||||
/// of the projected input to `±(1 + 1/(L-1))`, and an FSQ layer constructed
|
||||
/// with `preserve_symmetry=True` and a hard-clamped bound. Symmetric FSQ puts
|
||||
/// level `k` of `L` at `2k/(L-1) - 1`, spanning `[-1, 1]`; the asymmetric
|
||||
/// formula (`(k - L/2) / (L/2)`) is not what the checkpoint was trained with.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ResidualFSQ {
|
||||
project_in: Linear,
|
||||
project_out: Linear,
|
||||
/// Precomputed: (levels - 1) * (1 + eps) / 2 for each codebook dim
|
||||
half_l: Vec<f32>,
|
||||
/// Precomputed: levels // 2
|
||||
half_width: Vec<f32>,
|
||||
/// Precomputed: 0.5 for even levels, 0.0 for odd
|
||||
offset: Vec<f32>,
|
||||
/// Precomputed: atanh(offset / half_l) shift
|
||||
shift: Vec<f32>,
|
||||
/// Precomputed basis for index calculation: [1, L0, L0*L1, ...]
|
||||
/// Basis for the flat index: `[1, L0, L0*L1, ...]`.
|
||||
basis: Vec<i64>,
|
||||
/// Quantization levels per codebook dimension (e.g. [8, 8, 8, 5, 5, 5]).
|
||||
levels: Vec<usize>,
|
||||
@@ -290,22 +286,6 @@ impl ResidualFSQ {
|
||||
let project_in = linear(dim, codebook_dim, vb.pp("project_in"))?;
|
||||
let project_out = linear(codebook_dim, dim, vb.pp("project_out"))?;
|
||||
|
||||
let eps = 1e-3f32;
|
||||
let half_l: Vec<f32> = levels
|
||||
.iter()
|
||||
.map(|&l| (l as f32 - 1.0) * (1.0 + eps) / 2.0)
|
||||
.collect();
|
||||
let half_width: Vec<f32> = levels.iter().map(|&l| (l / 2) as f32).collect();
|
||||
let offset: Vec<f32> = levels
|
||||
.iter()
|
||||
.map(|&l| if l % 2 == 0 { 0.5 } else { 0.0 })
|
||||
.collect();
|
||||
let shift: Vec<f32> = half_l
|
||||
.iter()
|
||||
.zip(&offset)
|
||||
.map(|(&hl, &o)| (o / hl).atanh())
|
||||
.collect();
|
||||
|
||||
let mut basis = vec![1i64];
|
||||
for i in 0..levels.len() - 1 {
|
||||
basis.push(basis[i] * levels[i] as i64);
|
||||
@@ -314,76 +294,54 @@ impl ResidualFSQ {
|
||||
Ok(Self {
|
||||
project_in,
|
||||
project_out,
|
||||
half_l,
|
||||
half_width,
|
||||
offset,
|
||||
shift,
|
||||
basis,
|
||||
levels: levels.to_vec(),
|
||||
model_dtype,
|
||||
})
|
||||
}
|
||||
|
||||
/// FSQ bound + round: `round(tanh(z + shift) * half_l - offset) / half_width`
|
||||
fn fsq_quantize(&self, z: &Tensor) -> Result<Tensor> {
|
||||
let dev = z.device();
|
||||
let half_l = Tensor::new(&self.half_l[..], dev)?
|
||||
.unsqueeze(0)?
|
||||
.unsqueeze(0)?;
|
||||
let offset = Tensor::new(&self.offset[..], dev)?
|
||||
.unsqueeze(0)?
|
||||
.unsqueeze(0)?;
|
||||
let shift = Tensor::new(&self.shift[..], dev)?
|
||||
.unsqueeze(0)?
|
||||
.unsqueeze(0)?;
|
||||
let half_width = Tensor::new(&self.half_width[..], dev)?
|
||||
.unsqueeze(0)?
|
||||
.unsqueeze(0)?;
|
||||
/// A per-codebook-dimension constant as a `(1, 1, d)` f32 tensor.
|
||||
fn per_dim(&self, dev: &candle_core::Device, f: impl Fn(f32) -> f32) -> Result<Tensor> {
|
||||
let v: Vec<f32> = self.levels.iter().map(|&l| f(l as f32)).collect();
|
||||
Tensor::new(&v[..], dev)?.unsqueeze(0)?.unsqueeze(0)
|
||||
}
|
||||
|
||||
// bound: tanh(z + shift) * half_l - offset, then round, then / half_width
|
||||
let bounded = z
|
||||
.broadcast_add(&shift)?
|
||||
.tanh()?
|
||||
.broadcast_mul(&half_l)?
|
||||
.broadcast_sub(&offset)?;
|
||||
let rounded = bounded.round()?;
|
||||
rounded.broadcast_div(&half_width)
|
||||
/// Soft clamp, then the symmetric hard-clamped bound:
|
||||
/// `2/(L-1) * floor((L-1) * (clamp(z, -1, 1) + 1) / 2 + 0.5) - 1`.
|
||||
/// Returns `(codes, level_indices)`.
|
||||
fn fsq_quantize(&self, z: &Tensor) -> Result<(Tensor, Tensor)> {
|
||||
let dev = z.device();
|
||||
let z = z.to_dtype(DType::F32)?;
|
||||
let clamp = self.per_dim(dev, |l| 1.0 + 1.0 / (l - 1.0))?;
|
||||
let z = z.broadcast_div(&clamp)?.tanh()?.broadcast_mul(&clamp)?;
|
||||
let lm1 = self.per_dim(dev, |l| l - 1.0)?;
|
||||
let levels = z
|
||||
.clamp(-1f32, 1f32)?
|
||||
.affine(0.5, 0.5)?
|
||||
.broadcast_mul(&lm1)?
|
||||
.affine(1.0, 0.5)?
|
||||
.floor()?;
|
||||
let codes = levels.broadcast_div(&lm1)?.affine(2.0, -1.0)?;
|
||||
Ok((codes, levels))
|
||||
}
|
||||
|
||||
/// Quantize continuous representations into discrete tokens.
|
||||
/// Input: `(B, N, dim)`. Returns `(quantized_out, indices)` where
|
||||
/// `quantized_out` has the same shape as input and `indices` is `(B, N, num_quantizers)`.
|
||||
pub fn forward(&self, x: &Tensor) -> Result<(Tensor, Tensor)> {
|
||||
// project_in: (B, N, dim) -> (B, N, codebook_dim)
|
||||
let projected = x.apply(&self.project_in)?;
|
||||
|
||||
// FSQ quantize
|
||||
let codes = self.fsq_quantize(&projected)?; // (B, N, codebook_dim)
|
||||
|
||||
// Compute indices from codes: scale_and_shift then dot with basis
|
||||
let dev = x.device();
|
||||
let half_width = Tensor::new(&self.half_width[..], dev)?
|
||||
let (codes, levels) = self.fsq_quantize(&projected)?;
|
||||
let basis: Vec<f32> = self.basis.iter().map(|&b| b as f32).collect();
|
||||
let basis = Tensor::new(&basis[..], x.device())?
|
||||
.unsqueeze(0)?
|
||||
.unsqueeze(0)?;
|
||||
let offset = Tensor::new(&self.offset[..], dev)?
|
||||
.unsqueeze(0)?
|
||||
.unsqueeze(0)?;
|
||||
let basis = Tensor::new(&self.basis[..], dev)?
|
||||
.unsqueeze(0)?
|
||||
.unsqueeze(0)?;
|
||||
// scale_and_shift: codes * half_width + offset
|
||||
let shifted = codes.broadcast_mul(&half_width)?.broadcast_add(&offset)?;
|
||||
let indices = shifted
|
||||
.to_dtype(DType::F32)?
|
||||
.broadcast_mul(&basis.to_dtype(DType::F32)?)?
|
||||
let indices = levels
|
||||
.broadcast_mul(&basis)?
|
||||
.sum(2)?
|
||||
.round()?
|
||||
.to_dtype(DType::I64)?
|
||||
.unsqueeze(2)?; // (B, N, 1)
|
||||
|
||||
// project_out: (B, N, codebook_dim) -> (B, N, dim)
|
||||
let quantized_out = codes.apply(&self.project_out)?;
|
||||
|
||||
let quantized_out = codes.to_dtype(self.model_dtype)?.apply(&self.project_out)?;
|
||||
Ok((quantized_out, indices))
|
||||
}
|
||||
|
||||
@@ -393,42 +351,30 @@ impl ResidualFSQ {
|
||||
codes.apply(&self.project_out)
|
||||
}
|
||||
|
||||
/// Reconstruct from packed indices.
|
||||
/// Input: `(B, N, 1)` → Output: `(B, N, dim)`.
|
||||
///
|
||||
/// Inverse of the forward path: decompose flat index into per-level values
|
||||
/// via basis, then invert scale_and_shift, then `project_out`.
|
||||
pub fn get_output_from_indices(&self, indices: &Tensor) -> Result<Tensor> {
|
||||
/// The codes for packed indices: level `k = (index / basis) % L` of each
|
||||
/// dimension, at `2k/(L-1) - 1`. Input `(B, N, 1)` → `(B, N, codebook_dim)`.
|
||||
pub fn codes_from_indices(&self, indices: &Tensor) -> Result<Tensor> {
|
||||
let dev = indices.device();
|
||||
let indices = indices.squeeze(D::Minus1)?.to_dtype(DType::F32)?;
|
||||
let codebook_dim = self.levels.len();
|
||||
|
||||
// Decompose: level_val[i] = floor(index / basis[i]) % levels[i]
|
||||
// F32 has enough precision for indices up to 64000 (max codebook size).
|
||||
let mut code_dims = Vec::with_capacity(codebook_dim);
|
||||
for i in 0..codebook_dim {
|
||||
let b = self.basis[i] as f32;
|
||||
let l = self.levels[i] as f32;
|
||||
let dim_val = indices.affine(1.0 / b as f64, 0.0)?.floor()?;
|
||||
// modulo: dim_val - floor(dim_val / l) * l
|
||||
let dim_val = (&dim_val
|
||||
- dim_val
|
||||
.affine(1.0 / l as f64, 0.0)?
|
||||
.floor()?
|
||||
.affine(l as f64, 0.0)?)?;
|
||||
code_dims.push(dim_val);
|
||||
// F32 holds every index exactly: the codebook has 64000 entries.
|
||||
let mut levels = Vec::with_capacity(self.levels.len());
|
||||
for (i, &l) in self.levels.iter().enumerate() {
|
||||
let l = l as f64;
|
||||
let v = indices.affine(1.0 / self.basis[i] as f64, 0.0)?.floor()?;
|
||||
let v = (&v - v.affine(1.0 / l, 0.0)?.floor()?.affine(l, 0.0)?)?;
|
||||
levels.push(v);
|
||||
}
|
||||
let shifted = Tensor::stack(&code_dims, D::Minus1)?;
|
||||
let levels = Tensor::stack(&levels, D::Minus1)?;
|
||||
let lm1 = self.per_dim(dev, |l| l - 1.0)?;
|
||||
levels.broadcast_div(&lm1)?.affine(2.0, -1.0)
|
||||
}
|
||||
|
||||
// Inverse scale_and_shift (Python convention):
|
||||
// forward: shifted = codes * half_width + half_width → shifted ∈ [0, levels-1]
|
||||
// inverse: codes = (shifted - half_width) / half_width
|
||||
let hw = Tensor::new(&self.half_width[..], dev)?
|
||||
.unsqueeze(0)?
|
||||
.unsqueeze(0)?;
|
||||
let codes = shifted.broadcast_sub(&hw)?.broadcast_div(&hw)?;
|
||||
|
||||
codes.to_dtype(self.model_dtype)?.apply(&self.project_out)
|
||||
/// Reconstruct from packed indices.
|
||||
/// Input: `(B, N, 1)` → Output: `(B, N, dim)`.
|
||||
pub fn get_output_from_indices(&self, indices: &Tensor) -> Result<Tensor> {
|
||||
self.codes_from_indices(indices)?
|
||||
.to_dtype(self.model_dtype)?
|
||||
.apply(&self.project_out)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -437,104 +383,93 @@ mod tests {
|
||||
use super::*;
|
||||
use candle_core::Device;
|
||||
|
||||
/// Build a minimal ResidualFSQ with identity projections for testing.
|
||||
/// A ResidualFSQ with identity projections, as the library builds one
|
||||
/// when `dim` equals the codebook dimension.
|
||||
fn test_fsq() -> Result<ResidualFSQ> {
|
||||
let dev = &Device::Cpu;
|
||||
let levels = vec![8, 8, 8, 5, 5, 5];
|
||||
let codebook_dim = levels.len();
|
||||
|
||||
// Identity-like project_in/project_out (codebook_dim x codebook_dim)
|
||||
let eye = Tensor::eye(codebook_dim, DType::F32, dev)?;
|
||||
let zeros = Tensor::zeros(codebook_dim, DType::F32, dev)?;
|
||||
|
||||
let eps = 1e-3f32;
|
||||
let half_l: Vec<f32> = levels
|
||||
.iter()
|
||||
.map(|&l| (l as f32 - 1.0) * (1.0 + eps) / 2.0)
|
||||
.collect();
|
||||
let half_width: Vec<f32> = levels.iter().map(|&l| (l / 2) as f32).collect();
|
||||
let offset: Vec<f32> = levels
|
||||
.iter()
|
||||
.map(|&l| if l % 2 == 0 { 0.5 } else { 0.0 })
|
||||
.collect();
|
||||
let shift: Vec<f32> = half_l
|
||||
.iter()
|
||||
.zip(&offset)
|
||||
.map(|(&hl, &o)| (o / hl).atanh())
|
||||
.collect();
|
||||
let mut basis = vec![1i64];
|
||||
for i in 0..levels.len() - 1 {
|
||||
basis.push(basis[i] * levels[i] as i64);
|
||||
}
|
||||
|
||||
Ok(ResidualFSQ {
|
||||
project_in: Linear::from_weights(eye.clone(), Some(zeros.clone())),
|
||||
project_out: Linear::from_weights(eye, Some(zeros)),
|
||||
half_l,
|
||||
half_width,
|
||||
offset,
|
||||
shift,
|
||||
basis,
|
||||
levels,
|
||||
model_dtype: DType::F32,
|
||||
})
|
||||
}
|
||||
|
||||
fn close(got: &[f32], want: &[f32]) {
|
||||
for (i, (g, w)) in got.iter().zip(want).enumerate() {
|
||||
assert!((g - w).abs() < 1e-5, "dim {i}: got {g}, want {w}");
|
||||
}
|
||||
}
|
||||
|
||||
/// Recorded from vector_quantize_pytorch 1.27.20:
|
||||
/// `FSQ([8,8,8,5,5,5], preserve_symmetry=True).indices_to_codes(...)`.
|
||||
#[test]
|
||||
fn test_get_output_from_indices_roundtrip() -> Result<()> {
|
||||
// Test that manually-packed Python-convention indices correctly
|
||||
// reconstruct codes via get_output_from_indices.
|
||||
fn indices_decode_to_symmetric_codes() -> Result<()> {
|
||||
let fsq = test_fsq()?;
|
||||
let dev = &Device::Cpu;
|
||||
// levels = [8,8,8,5,5,5], half_width = [4,4,4,2,2,2]
|
||||
// basis = [1, 8, 64, 512, 2560, 12800]
|
||||
//
|
||||
// Choose codes: [0.25, -0.5, 0.75, 0.5, -1.0, 0.0]
|
||||
// Python scale_and_shift: shifted = codes * half_width + half_width
|
||||
// = [0.25*4+4, -0.5*4+4, 0.75*4+4, 0.5*2+2, -1.0*2+2, 0.0*2+2]
|
||||
// = [5, 2, 7, 3, 0, 2]
|
||||
// index = 5*1 + 2*8 + 7*64 + 3*512 + 0*2560 + 2*12800
|
||||
// = 5 + 16 + 448 + 1536 + 0 + 25600 = 27605
|
||||
let expected_codes: Vec<f32> = vec![0.25, -0.5, 0.75, 0.5, -1.0, 0.0];
|
||||
let index: i64 = 27605;
|
||||
|
||||
let indices = Tensor::new(&[[index]], dev)?.unsqueeze(2)?; // (1, 1, 1)
|
||||
let reconstructed = fsq.get_output_from_indices(&indices)?;
|
||||
let vals = reconstructed.squeeze(0)?.squeeze(0)?.to_vec1::<f32>()?;
|
||||
|
||||
for (i, (&got, &want)) in vals.iter().zip(expected_codes.iter()).enumerate() {
|
||||
assert!((got - want).abs() < 1e-5, "dim {i}: got={got}, want={want}");
|
||||
let cases: [(i64, [f32; 6]); 4] = [
|
||||
(0, [-1.0; 6]),
|
||||
(27605, [3.0 / 7.0, -3.0 / 7.0, 1.0, 0.5, -1.0, 0.0]),
|
||||
(32036, [1.0 / 7.0, 1.0 / 7.0, 1.0 / 7.0, 0.0, 0.0, 0.0]),
|
||||
(63999, [1.0; 6]),
|
||||
];
|
||||
for (index, want) in cases {
|
||||
let indices = Tensor::new(&[[index]], &Device::Cpu)?.unsqueeze(2)?;
|
||||
let got = fsq
|
||||
.get_output_from_indices(&indices)?
|
||||
.flatten_all()?
|
||||
.to_vec1::<f32>()?;
|
||||
close(&got, &want);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Recorded from vector_quantize_pytorch 1.27.20:
|
||||
/// `ResidualFSQ(dim=6, levels=[8,8,8,5,5,5], num_quantizers=1)(x)`.
|
||||
#[test]
|
||||
fn test_get_output_from_indices_known_values() -> Result<()> {
|
||||
fn forward_matches_the_library() -> Result<()> {
|
||||
let fsq = test_fsq()?;
|
||||
let dev = &Device::Cpu;
|
||||
|
||||
// Index 0 should decompose to level_vals = [0, 0, 0, 0, 0, 0]
|
||||
// Python convention: codes = (level_vals - half_width) / half_width
|
||||
// For levels [8,8,8,5,5,5]: half_width=[4,4,4,2,2,2]
|
||||
// codes = [(0-4)/4, (0-4)/4, (0-4)/4, (0-2)/2, (0-2)/2, (0-2)/2]
|
||||
// = [-1, -1, -1, -1, -1, -1]
|
||||
let indices = Tensor::new(&[[0i64]], dev)?.unsqueeze(2)?; // (1, 1, 1)
|
||||
let result = fsq.get_output_from_indices(&indices)?;
|
||||
let vals = result.squeeze(0)?.squeeze(0)?.to_vec1::<f32>()?;
|
||||
assert!((vals[0] - (-1.0)).abs() < 1e-5, "val[0]={}", vals[0]);
|
||||
assert!((vals[3] - (-1.0)).abs() < 1e-5, "val[3]={}", vals[3]);
|
||||
|
||||
// Index that corresponds to center: half_width values for each level
|
||||
// level_vals = [4, 4, 4, 2, 2, 2]
|
||||
// packed index = 4*1 + 4*8 + 4*64 + 2*512 + 2*2560 + 2*12800
|
||||
// = 4 + 32 + 256 + 1024 + 5120 + 25600 = 32036
|
||||
// codes = [0, 0, 0, 0, 0, 0]
|
||||
let indices = Tensor::new(&[[32036i64]], dev)?.unsqueeze(2)?;
|
||||
let result = fsq.get_output_from_indices(&indices)?;
|
||||
let vals = result.squeeze(0)?.squeeze(0)?.to_vec1::<f32>()?;
|
||||
for (i, &v) in vals.iter().enumerate() {
|
||||
assert!(v.abs() < 1e-5, "center val[{i}]={v}, expected 0");
|
||||
}
|
||||
let x = Tensor::new(
|
||||
&[[
|
||||
[0.3f32, -0.9, 2.0, 0.1, -0.05, 0.6],
|
||||
[-3.0, 0.0, 0.49, -0.6, 0.9, -0.2],
|
||||
]],
|
||||
&Device::Cpu,
|
||||
)?;
|
||||
let (q, idx) = fsq.forward(&x)?;
|
||||
let q = q.flatten_all()?.to_vec1::<f32>()?;
|
||||
close(
|
||||
&q,
|
||||
&[
|
||||
3.0 / 7.0,
|
||||
-5.0 / 7.0,
|
||||
1.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.5,
|
||||
-1.0,
|
||||
1.0 / 7.0,
|
||||
3.0 / 7.0,
|
||||
-0.5,
|
||||
1.0,
|
||||
0.0,
|
||||
],
|
||||
);
|
||||
assert_eq!(idx.flatten_all()?.to_vec1::<i64>()?, vec![45005, 36704]);
|
||||
// And the indices decode back to the same codes.
|
||||
let back = fsq
|
||||
.get_output_from_indices(&idx)?
|
||||
.flatten_all()?
|
||||
.to_vec1::<f32>()?;
|
||||
close(&back, &q);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -171,7 +171,10 @@ impl Module for OobleckEncoderBlock {
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct OobleckDecoderBlock {
|
||||
snake1: Snake1d,
|
||||
/// Built unpadded; the padding is trimmed from its output instead. See
|
||||
/// [`upsample`].
|
||||
conv_t1: ConvTranspose1d,
|
||||
trim: usize,
|
||||
res_unit1: OobleckResidualUnit,
|
||||
res_unit2: OobleckResidualUnit,
|
||||
res_unit3: OobleckResidualUnit,
|
||||
@@ -182,7 +185,7 @@ impl OobleckDecoderBlock {
|
||||
let snake1 = Snake1d::new(input_dim, vb.pp("snake1"))?;
|
||||
let cfg = ConvTranspose1dConfig {
|
||||
stride,
|
||||
padding: stride.div_ceil(2),
|
||||
padding: 0,
|
||||
..Default::default()
|
||||
};
|
||||
let conv_t1 = encodec::conv_transpose1d_weight_norm(
|
||||
@@ -199,6 +202,7 @@ impl OobleckDecoderBlock {
|
||||
Ok(Self {
|
||||
snake1,
|
||||
conv_t1,
|
||||
trim: stride.div_ceil(2),
|
||||
res_unit1,
|
||||
res_unit2,
|
||||
res_unit3,
|
||||
@@ -208,14 +212,25 @@ impl OobleckDecoderBlock {
|
||||
|
||||
impl Module for OobleckDecoderBlock {
|
||||
fn forward(&self, xs: &Tensor) -> Result<Tensor> {
|
||||
xs.apply(&self.snake1)?
|
||||
.apply(&self.conv_t1)?
|
||||
upsample(&xs.apply(&self.snake1)?, &self.conv_t1, self.trim)?
|
||||
.apply(&self.res_unit1)?
|
||||
.apply(&self.res_unit2)?
|
||||
.apply(&self.res_unit3)
|
||||
}
|
||||
}
|
||||
|
||||
/// A transposed convolution with `padding = trim`, computed as the
|
||||
/// unpadded convolution with `trim` samples cut from each end, which is
|
||||
/// the same result. candle's CUDA backend only takes its matmul (col2im)
|
||||
/// path for unpadded transposed convolutions; with padding it falls back
|
||||
/// to a direct kernel that made the VAE decode some 30x slower than the
|
||||
/// denoise on an RTX 3060.
|
||||
fn upsample(xs: &Tensor, conv: &ConvTranspose1d, trim: usize) -> Result<Tensor> {
|
||||
let ys = xs.apply(conv)?;
|
||||
let len = ys.dim(2)?;
|
||||
ys.narrow(2, trim, len - 2 * trim)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// OobleckEncoder
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -628,8 +643,10 @@ mod tests {
|
||||
decoder_channels: 8,
|
||||
decoder_input_channels: 4,
|
||||
audio_channels: 2,
|
||||
// Keep the real upsampling structure to exercise all 5 blocks.
|
||||
downsampling_ratios: vec![2, 4, 4, 6, 10],
|
||||
// Five blocks, as in production, but 2x each (hop 32 instead of
|
||||
// 1920): the tiling logic is the same, and the full-size hop made
|
||||
// this test peak at 4.6 GB, which got the CI test binary killed.
|
||||
downsampling_ratios: vec![2, 2, 2, 2, 2],
|
||||
channel_multiples: vec![1, 1, 1, 1, 1],
|
||||
..VaeConfig::default()
|
||||
};
|
||||
@@ -716,4 +733,39 @@ mod tests {
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unpadded_then_trimmed_matches_the_padded_transposed_conv() -> Result<()> {
|
||||
let dev = Device::Cpu;
|
||||
for stride in [2usize, 4, 5, 8, 10] {
|
||||
let k = 2 * stride;
|
||||
let pad = stride.div_ceil(2);
|
||||
let x = Tensor::randn(0f32, 1.0, (1, 6, 9), &dev)?;
|
||||
let w = Tensor::randn(0f32, 1.0, (6, 4, k), &dev)?;
|
||||
let b = Tensor::randn(0f32, 1.0, 4, &dev)?;
|
||||
let padded = ConvTranspose1d::new(
|
||||
w.clone(),
|
||||
Some(b.clone()),
|
||||
ConvTranspose1dConfig {
|
||||
stride,
|
||||
padding: pad,
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
let unpadded = ConvTranspose1d::new(
|
||||
w,
|
||||
Some(b),
|
||||
ConvTranspose1dConfig {
|
||||
stride,
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
let want = x.apply(&padded)?;
|
||||
let got = upsample(&x, &unpadded, pad)?;
|
||||
assert_eq!(want.dims(), got.dims(), "stride {stride}");
|
||||
let diff = (want - got)?.abs()?.max_all()?.to_scalar::<f32>()?;
|
||||
assert!(diff < 1e-4, "stride {stride}: max diff {diff}");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
+1014
-95
File diff suppressed because it is too large
Load Diff
@@ -116,6 +116,8 @@ pub struct CandleHarness {
|
||||
/// Image-pipeline settings (#198): text-encoder residency and the
|
||||
/// resolution ceiling.
|
||||
image_cfg: crate::config::ImageConfig,
|
||||
/// `[harness.candle.audio]` (#385).
|
||||
audio_cfg: crate::config::AudioConfig,
|
||||
}
|
||||
|
||||
/// Devices/capabilities snapshot of a model entering auto-recovery
|
||||
@@ -2522,6 +2524,7 @@ impl CandleHarness {
|
||||
reasoning_budget_cfg: config.reasoning_budget.clone(),
|
||||
admission_cfg: config.admission.clone(),
|
||||
image_cfg: config.image.clone(),
|
||||
audio_cfg: config.audio.clone(),
|
||||
});
|
||||
// Background auto-recovery task (#17). Holds a `Weak` so it can't
|
||||
// keep the harness alive. Spawned only when a tokio runtime is
|
||||
@@ -4085,7 +4088,21 @@ impl Harness for CandleHarness {
|
||||
// ACE-Step music models (#383) take the audio path: one device,
|
||||
// unquantized (preflight rejected tp>1 and quant).
|
||||
if matches!(plan.format, super::preflight::SourceFormat::AceStep) {
|
||||
return self.load_audio_model(spec, &source_id).await;
|
||||
let files =
|
||||
super::audio::AudioFiles::AceStep(self.resolve_audio_files(&source_id).await?);
|
||||
return self
|
||||
.load_audio_model(spec, files, super::audio::ACESTEP_LIMITS)
|
||||
.await;
|
||||
}
|
||||
// HeartMuLa (#393): the LM repo, plus its codec and generation
|
||||
// config from sibling repos.
|
||||
if matches!(plan.format, super::preflight::SourceFormat::HeartMuLa) {
|
||||
let files = super::audio::AudioFiles::HeartMuLa(
|
||||
self.resolve_heartmula_files(&source_id).await?,
|
||||
);
|
||||
return self
|
||||
.load_audio_model(spec, files, super::audio::HEARTMULA_LIMITS)
|
||||
.await;
|
||||
}
|
||||
|
||||
let tp_size = spec.tensor_parallel.unwrap_or(1);
|
||||
@@ -4650,7 +4667,8 @@ impl CandleHarness {
|
||||
async fn load_audio_model(
|
||||
&self,
|
||||
spec: &ModelSpec,
|
||||
source_id: &cortex_core::source::ModelSourceId,
|
||||
files: super::audio::AudioFiles,
|
||||
limits: cortex_core::audio::AudioLimits,
|
||||
) -> Result<()> {
|
||||
let devices = spec.devices.clone().unwrap_or_else(|| vec![0]);
|
||||
let device = Self::pick_device(&devices)?;
|
||||
@@ -4665,8 +4683,22 @@ impl CandleHarness {
|
||||
),
|
||||
};
|
||||
|
||||
let files = self.resolve_audio_files(source_id).await?;
|
||||
let has_planner = files.planner.is_some();
|
||||
let has_planner =
|
||||
matches!(&files, super::audio::AudioFiles::AceStep(f) if f.planner.is_some());
|
||||
// A host may accept shorter songs than the model allows (#385), for
|
||||
// a card whose memory runs out first.
|
||||
let mut limits = limits;
|
||||
if let super::audio::AudioFiles::AceStep(_) = &files
|
||||
&& let Some(max) = self.audio_cfg.max_duration_secs
|
||||
{
|
||||
anyhow::ensure!(
|
||||
max >= limits.min_duration,
|
||||
"[harness.candle.audio] max_duration_secs {max} is below the model's minimum {}",
|
||||
limits.min_duration
|
||||
);
|
||||
limits.max_duration = limits.max_duration.min(max);
|
||||
limits.default_duration = limits.default_duration.min(limits.max_duration);
|
||||
}
|
||||
let handle = worker
|
||||
.load_audio(files, spec.model_id.clone())
|
||||
.await
|
||||
@@ -4685,7 +4717,7 @@ impl CandleHarness {
|
||||
worker,
|
||||
handle,
|
||||
has_planner,
|
||||
limits: super::audio::ACESTEP_LIMITS,
|
||||
limits,
|
||||
poisoned: Arc::new(AtomicBool::new(false)),
|
||||
admission: crate::harness::admission::AdmissionController::new(&admission_cfg),
|
||||
inference_lock: Arc::new(tokio::sync::Mutex::new(())),
|
||||
@@ -4726,8 +4758,38 @@ impl CandleHarness {
|
||||
.context("snapshot path shorter than its repo path")?
|
||||
.to_path_buf();
|
||||
}
|
||||
// `[harness.candle.audio] planner` names one planner, or none;
|
||||
// unset tries them largest first.
|
||||
let named = self.audio_cfg.planner.as_deref();
|
||||
let candidates: Vec<&str> = match named {
|
||||
Some(n) if n.eq_ignore_ascii_case("none") => Vec::new(),
|
||||
Some(n) => vec![n],
|
||||
None => super::audio::PLANNER_DIRS.to_vec(),
|
||||
};
|
||||
let mut planner = None;
|
||||
for dir in super::audio::PLANNER_DIRS {
|
||||
// A name with a slash is a repo of its own (the 0.6B planner ships
|
||||
// as `ACE-Step/acestep-5Hz-lm-0.6B`, not in the main snapshot),
|
||||
// fetched through the same registry; the engine joins an absolute
|
||||
// path onto the root unchanged.
|
||||
if let Some(id) = named.filter(|n| n.contains('/')) {
|
||||
let planner_repo = api.model(id.to_string());
|
||||
let mut dir = None;
|
||||
for file in ["config.json", "model.safetensors", "tokenizer.json"] {
|
||||
let path = planner_repo
|
||||
.get(file)
|
||||
.await
|
||||
.with_context(|| format!("fetch {file} from planner repo {id}"))?;
|
||||
dir = path.parent().map(|p| p.to_path_buf());
|
||||
}
|
||||
let dir = dir.context("planner repo resolved to no directory")?;
|
||||
planner = Some(dir.to_string_lossy().into_owned());
|
||||
}
|
||||
let candidates = if planner.is_some() {
|
||||
Vec::new()
|
||||
} else {
|
||||
candidates
|
||||
};
|
||||
for dir in candidates {
|
||||
let mut ok = true;
|
||||
for file in super::audio::planner_files(dir) {
|
||||
if let Err(e) = repo.get(&file).await {
|
||||
@@ -4741,7 +4803,72 @@ impl CandleHarness {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(super::audio::AceStepFiles { root, planner })
|
||||
if let (Some(n), None) = (named, &planner)
|
||||
&& !n.eq_ignore_ascii_case("none")
|
||||
{
|
||||
anyhow::bail!(
|
||||
"{display_id} has no planner '{n}' (named by [harness.candle.audio] planner)"
|
||||
);
|
||||
}
|
||||
let layout = super::audio::AceStepLayout::from_config(&self.audio_cfg)?;
|
||||
Ok(super::audio::AceStepFiles {
|
||||
root,
|
||||
planner,
|
||||
layout,
|
||||
})
|
||||
}
|
||||
|
||||
/// Resolve a HeartMuLa LM repo (#393) and the two sibling repos it
|
||||
/// needs — the codec and the generation config — through the same
|
||||
/// registry, downloading as needed.
|
||||
async fn resolve_heartmula_files(
|
||||
&self,
|
||||
source_id: &cortex_core::source::ModelSourceId,
|
||||
) -> Result<super::audio::HeartMuLaFiles> {
|
||||
let api = self.hf_api_for(&source_id.scheme)?;
|
||||
// Every file of a repo, fetched; returns the snapshot directory.
|
||||
async fn fetch_all(
|
||||
api: &hf_hub::api::tokio::Api,
|
||||
repo_path: String,
|
||||
) -> Result<std::path::PathBuf> {
|
||||
let repo = api.model(repo_path.clone());
|
||||
let info = repo
|
||||
.info()
|
||||
.await
|
||||
.with_context(|| format!("list files of {repo_path}"))?;
|
||||
let mut dir = None;
|
||||
for s in info.siblings {
|
||||
let path = repo
|
||||
.get(&s.rfilename)
|
||||
.await
|
||||
.with_context(|| format!("fetch {} from {repo_path}", s.rfilename))?;
|
||||
if s.rfilename == "config.json" || dir.is_none() && !s.rfilename.contains('/') {
|
||||
dir = path.parent().map(|p| p.to_path_buf());
|
||||
}
|
||||
}
|
||||
dir.with_context(|| format!("{repo_path} has no top-level files"))
|
||||
}
|
||||
let lm_dir = fetch_all(&api, source_id.repo_path()).await?;
|
||||
let codec_dir = fetch_all(&api, super::audio::HEARTCODEC_REPO.to_string()).await?;
|
||||
let gen_repo = api.model(super::audio::HEARTMULA_GEN_REPO.to_string());
|
||||
let tokenizer = gen_repo.get("tokenizer.json").await.with_context(|| {
|
||||
format!(
|
||||
"fetch tokenizer.json from {}",
|
||||
super::audio::HEARTMULA_GEN_REPO
|
||||
)
|
||||
})?;
|
||||
let gen_config = gen_repo.get("gen_config.json").await.with_context(|| {
|
||||
format!(
|
||||
"fetch gen_config.json from {}",
|
||||
super::audio::HEARTMULA_GEN_REPO
|
||||
)
|
||||
})?;
|
||||
Ok(super::audio::HeartMuLaFiles {
|
||||
lm_dir,
|
||||
codec_dir,
|
||||
tokenizer,
|
||||
gen_config,
|
||||
})
|
||||
}
|
||||
|
||||
/// Load a Laya decision model (#336): resolve the checkpoints' files,
|
||||
@@ -5168,7 +5295,7 @@ impl CandleHarness {
|
||||
|
||||
/// Render one audio generation (#382). The cortex job manager sends one
|
||||
/// of these per job and waits for the whole render; the reply is the
|
||||
/// WAV bytes plus the result metadata cortex records on the job.
|
||||
/// FLAC bytes plus the result metadata cortex records on the job.
|
||||
///
|
||||
/// Validated against the model's limits before admission; the render
|
||||
/// itself is one device-worker job. A request without a seed gets a
|
||||
@@ -5203,16 +5330,17 @@ impl CandleHarness {
|
||||
request
|
||||
.validate_limits(&audio.limits)
|
||||
.map_err(|detail| InferenceError::InvalidAudioParams { detail })?;
|
||||
// The planner runs only when asked for until its port is proven
|
||||
// against the reference (#384); `auto` means off until then.
|
||||
// `auto` (or nothing) follows the vendor default: plan the song
|
||||
// when the model has a planner. The port matches upstream's trace
|
||||
// (#384).
|
||||
let planner = match request.planner.as_deref() {
|
||||
Some("on") if !audio.has_planner => {
|
||||
return Err(InferenceError::InvalidAudioParams {
|
||||
detail: "this model has no song planner; use planner=off or auto".into(),
|
||||
});
|
||||
}
|
||||
Some("on") => true,
|
||||
_ => false,
|
||||
Some("off") => false,
|
||||
_ => audio.has_planner,
|
||||
};
|
||||
// A random seed small enough to type back in, so any render can be
|
||||
// reproduced from what the job reports.
|
||||
@@ -5235,6 +5363,7 @@ impl CandleHarness {
|
||||
planner,
|
||||
temperature: request.temperature,
|
||||
language: request.language.clone(),
|
||||
top_k: request.top_k.map(|k| k as usize),
|
||||
};
|
||||
|
||||
let _permit = audio.admission.enter(principal).await?;
|
||||
@@ -5248,13 +5377,14 @@ impl CandleHarness {
|
||||
let duration = result.duration_secs();
|
||||
let sample_rate = result.sample_rate;
|
||||
let timing = result.timing;
|
||||
let wav = tokio::task::spawn_blocking(move || {
|
||||
super::audio::encode_wav_pcm16(&result.left, &result.right, result.sample_rate)
|
||||
let flac = tokio::task::spawn_blocking(move || {
|
||||
super::audio::encode_flac_pcm16(&result.left, &result.right, result.sample_rate)
|
||||
})
|
||||
.await
|
||||
.map_err(|e| InferenceError::Other(anyhow::anyhow!("wav encode task: {e:#}")))?;
|
||||
.map_err(|e| InferenceError::Other(anyhow::anyhow!("flac encode task: {e:#}")))?
|
||||
.map_err(InferenceError::Other)?;
|
||||
Ok((
|
||||
wav,
|
||||
flac,
|
||||
cortex_core::audio::NeuronAudioResult {
|
||||
duration,
|
||||
sample_rate,
|
||||
|
||||
@@ -66,9 +66,18 @@ pub fn classify_files(
|
||||
// entry) is not servable, so claiming "text" would be a lie of
|
||||
// the same kind this module exists to remove.
|
||||
SourceFormat::Empty => Vec::new(),
|
||||
SourceFormat::HeartMuLa => vec![cortex_core::audio::CAPABILITY.to_string()],
|
||||
_ => {
|
||||
let config = read_config();
|
||||
// HeartMuLa (#393) is only recognisable by its config.
|
||||
if config
|
||||
.as_deref()
|
||||
.is_some_and(crate::harness::audio::is_heartmula_config)
|
||||
{
|
||||
return vec![cortex_core::audio::CAPABILITY.to_string()];
|
||||
}
|
||||
let mut caps = vec!["text".to_string()];
|
||||
if read_config().is_some_and(|text| declares_vision(&text)) {
|
||||
if config.is_some_and(|text| declares_vision(&text)) {
|
||||
caps.push("vision".to_string());
|
||||
}
|
||||
caps
|
||||
@@ -182,4 +191,19 @@ mod tests {
|
||||
let refs: Vec<&str> = files.iter().map(String::as_str).collect();
|
||||
assert_eq!(classify_files(&refs, || None), vec!["audio".to_string()]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn heartmula_snapshot_is_audio_by_its_config() {
|
||||
let files = [
|
||||
"config.json",
|
||||
"model-00001-of-00004.safetensors",
|
||||
"model.safetensors.index.json",
|
||||
];
|
||||
let caps = classify_files(&files, || {
|
||||
Some(r#"{"architectures": ["HeartMuLa"], "model_type": "heartmula"}"#.to_string())
|
||||
});
|
||||
assert_eq!(caps, vec!["audio".to_string()]);
|
||||
let text = classify_files(&files, || Some(r#"{"model_type": "qwen3"}"#.to_string()));
|
||||
assert_eq!(text, vec!["text".to_string()]);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -73,7 +73,7 @@ struct DeviceWorkerState {
|
||||
next_decision_handle: u64,
|
||||
/// ACE-Step music pipeline slab (#383). Same lifecycle discipline as
|
||||
/// `image_models`, in its own namespace.
|
||||
audio_models: HashMap<AudioHandle, Box<crate::harness::audio::AceStepEngine>>,
|
||||
audio_models: HashMap<AudioHandle, Box<crate::harness::audio::AudioEngine>>,
|
||||
/// Counter for minting fresh `AudioHandle`s.
|
||||
next_audio_handle: u64,
|
||||
/// Leader's NCCL state. Populated by `Job::NcclInit`; the
|
||||
@@ -823,14 +823,15 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
|
||||
model_id,
|
||||
reply,
|
||||
} => {
|
||||
let result = crate::harness::audio::AceStepEngine::load(&files, &state.device)
|
||||
.with_context(|| format!("load ace-step pipeline for {model_id}"))
|
||||
let result = crate::harness::audio::AudioEngine::load(&files, &state.device)
|
||||
.with_context(|| format!("load audio engine for {model_id}"))
|
||||
.map(|engine| {
|
||||
let handle = AudioHandle(state.next_audio_handle);
|
||||
state.next_audio_handle = state.next_audio_handle.wrapping_add(1);
|
||||
tracing::info!(
|
||||
device_index,
|
||||
model_id,
|
||||
engine = engine.name(),
|
||||
planner = engine.has_planner(),
|
||||
handle = handle.0,
|
||||
"device worker: audio pipeline loaded"
|
||||
|
||||
@@ -68,12 +68,12 @@ pub struct ImageHandle(pub u64);
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub struct DecisionHandle(pub u64);
|
||||
|
||||
/// Opaque handle to an ACE-Step music pipeline ([`AceStepEngine`]) stored
|
||||
/// Opaque handle to a music engine ([`AudioEngine`]) stored
|
||||
/// in the worker thread's audio slab (#383). Same discipline as
|
||||
/// [`ImageHandle`]: construction, every generation and `Drop` happen on the
|
||||
/// worker thread, and `Job::DropAudio` is the only way to release it.
|
||||
///
|
||||
/// [`AceStepEngine`]: crate::harness::audio::AceStepEngine
|
||||
/// [`AudioEngine`]: crate::harness::audio::AudioEngine
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub struct AudioHandle(pub u64);
|
||||
|
||||
@@ -597,7 +597,7 @@ pub enum Job {
|
||||
/// Load an ACE-Step snapshot (#383) — already on local disk — into
|
||||
/// the audio slab. File resolution (hf-hub) happens on the async side.
|
||||
LoadAudio {
|
||||
files: crate::harness::audio::AceStepFiles,
|
||||
files: crate::harness::audio::AudioFiles,
|
||||
model_id: String,
|
||||
reply: oneshot::Sender<Result<AudioHandle>>,
|
||||
},
|
||||
|
||||
@@ -339,10 +339,10 @@ impl DeviceWorkerHandle {
|
||||
}
|
||||
}
|
||||
|
||||
/// Load an ACE-Step music pipeline (#383) on the worker thread.
|
||||
/// Load a music engine (#383, #393) on the worker thread.
|
||||
pub async fn load_audio(
|
||||
&self,
|
||||
files: crate::harness::audio::AceStepFiles,
|
||||
files: crate::harness::audio::AudioFiles,
|
||||
model_id: String,
|
||||
) -> Result<AudioHandle, WorkerError> {
|
||||
if self.poisoned.load(Ordering::Acquire) {
|
||||
|
||||
@@ -59,6 +59,10 @@ pub enum SourceFormat {
|
||||
/// `config.json`. Served by the decision path — single device,
|
||||
/// unquantized. Takes precedence over the safetensors sniff.
|
||||
DecisionModel,
|
||||
/// A HeartMuLa LM repo (#393): sharded safetensors like any dense LLM,
|
||||
/// told apart by `model_type: heartmula` in its `config.json`. Served by
|
||||
/// the audio path with its codec from a sibling repo; single device.
|
||||
HeartMuLa,
|
||||
/// An ACE-Step 1.5 music repo (#383): the turbo DiT, text encoder,
|
||||
/// VAE and planner as subfolders of one snapshot. Served by the audio
|
||||
/// path — single device, unquantized for now. Takes precedence over
|
||||
@@ -181,7 +185,17 @@ pub async fn preflight(
|
||||
};
|
||||
|
||||
let filenames: Vec<&str> = owned_filenames.iter().map(String::as_str).collect();
|
||||
let format = classify(&filenames);
|
||||
let mut format = classify(&filenames);
|
||||
// A HeartMuLa LM (#393) looks like any dense checkpoint by its file
|
||||
// names; only its config says otherwise.
|
||||
if matches!(format, SourceFormat::DenseSafetensors { .. })
|
||||
&& filenames.contains(&"config.json")
|
||||
&& let Ok(path) = repo.get("config.json").await
|
||||
&& std::fs::read_to_string(&path)
|
||||
.is_ok_and(|t| crate::harness::audio::is_heartmula_config(&t))
|
||||
{
|
||||
format = SourceFormat::HeartMuLa;
|
||||
}
|
||||
let tp_size = spec.tensor_parallel.unwrap_or(1);
|
||||
|
||||
match (&format, tp_size, spec.quant.as_deref()) {
|
||||
@@ -264,17 +278,21 @@ pub async fn preflight(
|
||||
tp_size,
|
||||
picked_quant_file: None,
|
||||
}),
|
||||
(SourceFormat::AceStep, tp, _) if tp > 1 => Err(PreflightError::Unsupported {
|
||||
model_id: source_id.to_string(),
|
||||
reason: format!(
|
||||
"music models load on a single device; tensor_parallel={tp} is not supported"
|
||||
),
|
||||
}),
|
||||
(SourceFormat::AceStep, _, Some(q)) => Err(PreflightError::Unsupported {
|
||||
model_id: source_id.to_string(),
|
||||
reason: format!("music models load unquantized for now; remove quant={q:?}"),
|
||||
}),
|
||||
(SourceFormat::AceStep, _, None) => Ok(PlacementPlan {
|
||||
(SourceFormat::AceStep | SourceFormat::HeartMuLa, tp, _) if tp > 1 => {
|
||||
Err(PreflightError::Unsupported {
|
||||
model_id: source_id.to_string(),
|
||||
reason: format!(
|
||||
"music models load on a single device; tensor_parallel={tp} is not supported"
|
||||
),
|
||||
})
|
||||
}
|
||||
(SourceFormat::AceStep | SourceFormat::HeartMuLa, _, Some(q)) => {
|
||||
Err(PreflightError::Unsupported {
|
||||
model_id: source_id.to_string(),
|
||||
reason: format!("music models load unquantized for now; remove quant={q:?}"),
|
||||
})
|
||||
}
|
||||
(SourceFormat::AceStep | SourceFormat::HeartMuLa, _, None) => Ok(PlacementPlan {
|
||||
model_id: source_id.to_string(),
|
||||
format: format.clone(),
|
||||
tp_size,
|
||||
@@ -732,15 +750,19 @@ mod tests {
|
||||
tp_size,
|
||||
picked_quant_file: None,
|
||||
}),
|
||||
(SourceFormat::AceStep, tp, _) if tp > 1 => Err(PreflightError::Unsupported {
|
||||
model_id: source_id.to_string(),
|
||||
reason: format!("music models load on a single device; tensor_parallel={tp}"),
|
||||
}),
|
||||
(SourceFormat::AceStep, _, Some(q)) => Err(PreflightError::Unsupported {
|
||||
model_id: source_id.to_string(),
|
||||
reason: format!("music models load unquantized for now; remove quant={q:?}"),
|
||||
}),
|
||||
(SourceFormat::AceStep, _, None) => Ok(PlacementPlan {
|
||||
(SourceFormat::AceStep | SourceFormat::HeartMuLa, tp, _) if tp > 1 => {
|
||||
Err(PreflightError::Unsupported {
|
||||
model_id: source_id.to_string(),
|
||||
reason: format!("music models load on a single device; tensor_parallel={tp}"),
|
||||
})
|
||||
}
|
||||
(SourceFormat::AceStep | SourceFormat::HeartMuLa, _, Some(q)) => {
|
||||
Err(PreflightError::Unsupported {
|
||||
model_id: source_id.to_string(),
|
||||
reason: format!("music models load unquantized for now; remove quant={q:?}"),
|
||||
})
|
||||
}
|
||||
(SourceFormat::AceStep | SourceFormat::HeartMuLa, _, None) => Ok(PlacementPlan {
|
||||
model_id: source_id.to_string(),
|
||||
format: format.clone(),
|
||||
tp_size,
|
||||
|
||||
@@ -872,3 +872,76 @@ async fn test_audio_invalid_params_400() {
|
||||
assert_eq!(body["error"]["code"], "invalid_audio_params", "{bad}");
|
||||
}
|
||||
}
|
||||
|
||||
/// A harness whose load takes a while and records when it finishes.
|
||||
struct SlowLoadHarness {
|
||||
finished: Arc<std::sync::atomic::AtomicBool>,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl cortex_core::harness::Harness for SlowLoadHarness {
|
||||
fn name(&self) -> &str {
|
||||
"slow"
|
||||
}
|
||||
async fn health(&self) -> cortex_core::harness::HarnessHealth {
|
||||
cortex_core::harness::HarnessHealth {
|
||||
name: "slow".into(),
|
||||
running: true,
|
||||
uptime_secs: None,
|
||||
}
|
||||
}
|
||||
async fn list_models(&self) -> anyhow::Result<Vec<cortex_core::harness::ModelInfo>> {
|
||||
Ok(vec![])
|
||||
}
|
||||
async fn load_model(&self, _spec: &cortex_core::harness::ModelSpec) -> anyhow::Result<()> {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(400)).await;
|
||||
self.finished
|
||||
.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
Ok(())
|
||||
}
|
||||
async fn unload_model(&self, _model_id: &str) -> anyhow::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
async fn inference_endpoint(&self, _model_id: &str) -> Option<String> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// A client that gives up mid-load (a timeout during a long first fetch)
|
||||
/// must not cancel the load: it runs to completion regardless.
|
||||
#[tokio::test]
|
||||
async fn test_load_survives_client_disconnect() {
|
||||
let finished = Arc::new(std::sync::atomic::AtomicBool::new(false));
|
||||
let mut registry = HarnessRegistry::new();
|
||||
registry.register(Arc::new(SlowLoadHarness {
|
||||
finished: Arc::clone(&finished),
|
||||
}));
|
||||
let state = Arc::new(NeuronState {
|
||||
discovery: fake_discovery(),
|
||||
health_cache: Arc::new(HealthCache::new()),
|
||||
registry: RwLock::new(registry),
|
||||
candle: None,
|
||||
activation: Arc::new(ActivationTracker::new(&[])),
|
||||
});
|
||||
let app = api::neuron_routes().with_state(state);
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!("http://{addr}/models/load"))
|
||||
.timeout(std::time::Duration::from_millis(100))
|
||||
.json(&json!({"model_id": "some/model", "harness": "slow"}))
|
||||
.send()
|
||||
.await;
|
||||
assert!(resp.is_err(), "the client should have timed out mid-load");
|
||||
assert!(!finished.load(std::sync::atomic::Ordering::SeqCst));
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(800)).await;
|
||||
assert!(
|
||||
finished.load(std::sync::atomic::Ordering::SeqCst),
|
||||
"the load was cancelled with the request"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
---
|
||||
title: Using decisions
|
||||
sidebar_label: Decisions
|
||||
description: Typed questions over text or JSON, answered with calibrated probabilities — the System-One endpoint and its playground.
|
||||
description: Typed questions over text or JSON, answered with calibrated probabilities — the System-One endpoint and the Decisions page.
|
||||
---
|
||||
|
||||
# Using decisions
|
||||
@@ -13,7 +13,7 @@ request. That makes it fast (tens of milliseconds), cheap, and
|
||||
predictable enough to put in front of a workflow: routing a ticket,
|
||||
screening a message, deciding whether to escalate.
|
||||
|
||||
Try it without writing code in [the playground](/playground).
|
||||
Try it without writing code on [the Decisions page](/decisions).
|
||||
|
||||
The endpoint is `POST https://helexa.ai/v1/systemone`, and it speaks the
|
||||
TypeSafe Jev protocol, so existing Jev clients work unchanged.
|
||||
|
||||
+369
-16
@@ -1888,53 +1888,209 @@ html[dir="rtl"] .hx-landing-composer textarea {
|
||||
/* -- Music (/music) -------------------------------------------------------- */
|
||||
|
||||
/* The music page reuses the image form controls; these cover what images
|
||||
* has no counterpart for: labelled text areas and a job list. */
|
||||
* has no counterpart for. */
|
||||
.hx-music-field {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 0.25rem;
|
||||
gap: 0.35rem;
|
||||
min-width: 0;
|
||||
font-size: 0.82rem;
|
||||
color: var(--color-text-muted);
|
||||
}
|
||||
.hx-music-field textarea,
|
||||
.hx-music input[type="number"] {
|
||||
.hx-music-field textarea {
|
||||
border-radius: 0.5rem;
|
||||
border: 1px solid var(--color-border-subtle);
|
||||
background-color: var(--color-bg-elevated);
|
||||
color: var(--color-text);
|
||||
padding: 0.45rem 0.6rem;
|
||||
padding: 0.55rem 0.7rem;
|
||||
font: inherit;
|
||||
font-size: 0.9rem;
|
||||
font-size: 0.92rem;
|
||||
line-height: 1.5;
|
||||
min-height: 36px;
|
||||
resize: vertical;
|
||||
}
|
||||
.hx-music input[type="number"] {
|
||||
width: 7rem;
|
||||
/* Style beside lyrics on a wide screen: lyrics are short lines, so a
|
||||
* full-width box wastes most of its width. Both boxes share one height. */
|
||||
.hx-music-words {
|
||||
display: grid;
|
||||
gap: 1rem;
|
||||
}
|
||||
.hx-music-jobs {
|
||||
@media (min-width: 768px) {
|
||||
.hx-music-words.two {
|
||||
grid-template-columns: minmax(0, 1fr) minmax(0, 1fr);
|
||||
}
|
||||
.hx-music-words.two .hx-music-field textarea {
|
||||
flex: 1 1 auto;
|
||||
}
|
||||
}
|
||||
.hx-music-tags {
|
||||
display: flex;
|
||||
flex-wrap: wrap;
|
||||
gap: 0.35rem;
|
||||
}
|
||||
.hx-music-tag {
|
||||
padding: 0.2rem 0.65rem;
|
||||
border: 1px solid var(--color-border-subtle);
|
||||
border-radius: 999px;
|
||||
background: none;
|
||||
color: var(--color-text-muted);
|
||||
font-size: 0.78rem;
|
||||
transition:
|
||||
border-color var(--transition-fast),
|
||||
color var(--transition-fast),
|
||||
background-color var(--transition-fast);
|
||||
}
|
||||
.hx-music-tag:hover:not(:disabled) {
|
||||
border-color: var(--color-accent);
|
||||
color: var(--color-text);
|
||||
}
|
||||
.hx-music-tag.on {
|
||||
border-color: var(--color-accent);
|
||||
background-color: var(--color-accent-soft);
|
||||
color: var(--color-text);
|
||||
}
|
||||
.hx-music-length {
|
||||
flex: 1 1 14rem;
|
||||
max-width: 22rem;
|
||||
}
|
||||
.hx-music-length strong {
|
||||
color: var(--color-text);
|
||||
font-variant-numeric: tabular-nums;
|
||||
}
|
||||
.hx-music-length input[type="range"] {
|
||||
width: 100%;
|
||||
accent-color: var(--color-accent);
|
||||
}
|
||||
/* A disabled generate button must read as unavailable at a glance, not
|
||||
* as a slightly faded gradient. */
|
||||
.hx-music .hx-music-generate:disabled {
|
||||
background: var(--color-bg-subtle);
|
||||
color: var(--color-text-muted);
|
||||
box-shadow: inset 0 0 0 1px var(--color-border-subtle);
|
||||
opacity: 1;
|
||||
cursor: not-allowed;
|
||||
}
|
||||
/* Phone widths: the history button sits over the title's start. */
|
||||
@media (max-width: 768px) {
|
||||
.hx-music > h1 {
|
||||
padding-inline-start: 2.75rem;
|
||||
}
|
||||
}
|
||||
.hx-music-blocker {
|
||||
margin: -0.4rem 0 0;
|
||||
text-align: end;
|
||||
color: var(--color-text-muted);
|
||||
}
|
||||
.hx-music-spin {
|
||||
animation: hx-music-turn 0.9s linear infinite;
|
||||
}
|
||||
@keyframes hx-music-turn {
|
||||
to {
|
||||
transform: rotate(360deg);
|
||||
}
|
||||
}
|
||||
|
||||
/* The studio: what is happening now, below the form. Empty on arrival. */
|
||||
.hx-music-studio {
|
||||
margin-top: 1.5rem;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 0.75rem;
|
||||
}
|
||||
.hx-music-job {
|
||||
padding: 0.7rem 0.8rem;
|
||||
border: 1px solid var(--color-border-subtle);
|
||||
.hx-music-studio:empty {
|
||||
display: none;
|
||||
}
|
||||
.hx-music-activity {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 0.75rem;
|
||||
padding: 0.8rem 1rem;
|
||||
border: 1px dashed var(--color-border-strong);
|
||||
border-radius: 0.75rem;
|
||||
color: var(--color-text-muted);
|
||||
font-size: 0.9rem;
|
||||
}
|
||||
.hx-music-job {
|
||||
display: grid;
|
||||
grid-template-columns: 5.5rem minmax(0, 1fr);
|
||||
gap: 1rem;
|
||||
padding: 0.9rem;
|
||||
border: 1px solid var(--color-border-subtle);
|
||||
border-radius: 0.9rem;
|
||||
background-color: var(--color-bg-elevated);
|
||||
}
|
||||
.hx-music-cover {
|
||||
width: 5.5rem;
|
||||
aspect-ratio: 1;
|
||||
border-radius: 0.6rem;
|
||||
display: grid;
|
||||
place-items: center;
|
||||
position: relative;
|
||||
overflow: hidden;
|
||||
box-shadow: 0 6px 18px rgb(0 0 0 / 0.25);
|
||||
}
|
||||
/* A record's label and groove sheen over the gradient. */
|
||||
.hx-music-cover::after {
|
||||
content: "";
|
||||
position: absolute;
|
||||
inset: 18%;
|
||||
border-radius: 50%;
|
||||
background: radial-gradient(circle, rgb(255 255 255 / 0.35) 0 12%, transparent 13% 100%),
|
||||
repeating-radial-gradient(circle, rgb(0 0 0 / 0.12) 0 1px, transparent 1px 4px);
|
||||
mix-blend-mode: overlay;
|
||||
}
|
||||
.hx-music-cover .hx-music-eq {
|
||||
position: relative;
|
||||
z-index: 1;
|
||||
color: #fff;
|
||||
height: 1.6rem;
|
||||
}
|
||||
.hx-music-job-body {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 0.4rem;
|
||||
min-width: 0;
|
||||
}
|
||||
.hx-music-job header {
|
||||
display: flex;
|
||||
align-items: flex-start;
|
||||
gap: 0.5rem;
|
||||
}
|
||||
.hx-music-title {
|
||||
margin: 0;
|
||||
flex: 1 1 auto;
|
||||
min-width: 0;
|
||||
font-size: 1rem;
|
||||
font-weight: 600;
|
||||
display: -webkit-box;
|
||||
-webkit-line-clamp: 2;
|
||||
-webkit-box-orient: vertical;
|
||||
overflow: hidden;
|
||||
}
|
||||
.hx-music-close {
|
||||
flex-shrink: 0;
|
||||
margin-block-start: -0.25rem;
|
||||
}
|
||||
.hx-music-prompt {
|
||||
margin: 0;
|
||||
font-size: 0.85rem;
|
||||
color: var(--color-text-muted);
|
||||
display: -webkit-box;
|
||||
-webkit-line-clamp: 2;
|
||||
-webkit-box-orient: vertical;
|
||||
overflow: hidden;
|
||||
}
|
||||
.hx-music-meta {
|
||||
display: flex;
|
||||
flex-wrap: wrap;
|
||||
gap: 0.6rem;
|
||||
align-items: baseline;
|
||||
gap: 0.25rem 0.75rem;
|
||||
font-size: 0.8rem;
|
||||
color: var(--color-text-muted);
|
||||
}
|
||||
.hx-music-status {
|
||||
font-size: 0.8rem;
|
||||
font-weight: 600;
|
||||
color: var(--color-text);
|
||||
}
|
||||
.hx-music-running .hx-music-status,
|
||||
.hx-music-queued .hx-music-status {
|
||||
@@ -1951,9 +2107,206 @@ html[dir="rtl"] .hx-landing-composer textarea {
|
||||
.hx-music-job .hx-image-card-actions {
|
||||
flex-wrap: wrap;
|
||||
}
|
||||
@media (max-width: 575.98px) {
|
||||
.hx-music-job {
|
||||
grid-template-columns: 3.5rem minmax(0, 1fr);
|
||||
gap: 0.75rem;
|
||||
}
|
||||
.hx-music-cover {
|
||||
width: 3.5rem;
|
||||
}
|
||||
}
|
||||
|
||||
/* Equaliser bars: the "something is being made" signal. */
|
||||
.hx-music-eq {
|
||||
display: inline-flex;
|
||||
align-items: flex-end;
|
||||
gap: 3px;
|
||||
height: 1.1rem;
|
||||
color: var(--color-accent);
|
||||
}
|
||||
.hx-music-eq i {
|
||||
width: 3px;
|
||||
height: 100%;
|
||||
border-radius: 2px;
|
||||
background: currentColor;
|
||||
transform-origin: bottom;
|
||||
animation: hx-music-bounce 0.9s ease-in-out infinite;
|
||||
}
|
||||
.hx-music-eq i:nth-child(2) {
|
||||
animation-delay: -0.3s;
|
||||
}
|
||||
.hx-music-eq i:nth-child(3) {
|
||||
animation-delay: -0.6s;
|
||||
}
|
||||
.hx-music-eq i:nth-child(4) {
|
||||
animation-delay: -0.15s;
|
||||
}
|
||||
@keyframes hx-music-bounce {
|
||||
0%,
|
||||
100% {
|
||||
transform: scaleY(0.25);
|
||||
}
|
||||
50% {
|
||||
transform: scaleY(1);
|
||||
}
|
||||
}
|
||||
|
||||
/* The player: a round play button, a waveform that is the seek bar, the
|
||||
* time. */
|
||||
.hx-music-player {
|
||||
width: min(100%, 24rem);
|
||||
height: 36px;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 0.75rem;
|
||||
margin-block: 0.25rem;
|
||||
}
|
||||
.hx-music-play {
|
||||
flex-shrink: 0;
|
||||
width: 2.6rem;
|
||||
height: 2.6rem;
|
||||
border: none;
|
||||
border-radius: 50%;
|
||||
display: grid;
|
||||
place-items: center;
|
||||
background: var(--gradient-helix);
|
||||
color: #fff;
|
||||
transition: transform var(--transition-fast);
|
||||
}
|
||||
.hx-music-play:hover:not(:disabled) {
|
||||
transform: scale(1.06);
|
||||
}
|
||||
html[dir="rtl"] .hx-music-play svg {
|
||||
transform: scaleX(-1);
|
||||
}
|
||||
.hx-music-wave {
|
||||
flex: 1 1 auto;
|
||||
min-width: 0;
|
||||
height: 2.6rem;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 2px;
|
||||
cursor: pointer;
|
||||
}
|
||||
.hx-music-wave i {
|
||||
flex: 1 1 0;
|
||||
min-width: 0;
|
||||
border-radius: 1px;
|
||||
background-color: var(--color-border-strong);
|
||||
transition: background-color 80ms linear;
|
||||
}
|
||||
.hx-music-wave i.on {
|
||||
background-color: var(--color-accent);
|
||||
}
|
||||
.hx-music-player.standin .hx-music-wave i {
|
||||
opacity: 0.55;
|
||||
}
|
||||
.hx-music-wave:focus-visible {
|
||||
outline: 2px solid var(--color-accent);
|
||||
outline-offset: 3px;
|
||||
border-radius: 4px;
|
||||
}
|
||||
.hx-music-time {
|
||||
flex-shrink: 0;
|
||||
font-size: 0.78rem;
|
||||
color: var(--color-text-muted);
|
||||
font-variant-numeric: tabular-nums;
|
||||
}
|
||||
.hx-music-spinner {
|
||||
width: 1rem;
|
||||
height: 1rem;
|
||||
border-radius: 50%;
|
||||
border: 2px solid rgb(255 255 255 / 0.4);
|
||||
border-top-color: #fff;
|
||||
animation: hx-music-turn 0.8s linear infinite;
|
||||
}
|
||||
@media (max-width: 575.98px) {
|
||||
.hx-music-time {
|
||||
display: none;
|
||||
}
|
||||
.hx-music-wave {
|
||||
gap: 1px;
|
||||
}
|
||||
}
|
||||
@media (prefers-reduced-motion: reduce) {
|
||||
.hx-music-eq i,
|
||||
.hx-music-spin,
|
||||
.hx-music-spinner {
|
||||
animation: none;
|
||||
}
|
||||
.hx-music-eq i {
|
||||
transform: scaleY(0.6);
|
||||
}
|
||||
}
|
||||
|
||||
/* History sidebar (#402): the chat sidebar, with collapsible sections. */
|
||||
.hx-group-toggle {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 0.35rem;
|
||||
min-width: 0;
|
||||
padding: 0;
|
||||
border: none;
|
||||
background: none;
|
||||
color: inherit;
|
||||
font: inherit;
|
||||
letter-spacing: inherit;
|
||||
text-transform: inherit;
|
||||
text-align: start;
|
||||
}
|
||||
.hx-group-count {
|
||||
margin-inline-start: auto;
|
||||
padding-inline: 0.4rem;
|
||||
font-weight: 500;
|
||||
opacity: 0.7;
|
||||
}
|
||||
html[dir="rtl"] .hx-flip-rtl {
|
||||
transform: scaleX(-1);
|
||||
}
|
||||
.hx-music-item {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 0.5rem;
|
||||
min-width: 0;
|
||||
}
|
||||
.hx-music-when {
|
||||
flex-shrink: 0;
|
||||
font-size: 0.75rem;
|
||||
opacity: 0.75;
|
||||
}
|
||||
.hx-music-dot {
|
||||
flex-shrink: 0;
|
||||
width: 0.45rem;
|
||||
height: 0.45rem;
|
||||
border-radius: 50%;
|
||||
background-color: var(--color-border-strong);
|
||||
}
|
||||
.hx-music-dot-succeeded {
|
||||
background-color: var(--color-accent);
|
||||
}
|
||||
.hx-music-dot-queued,
|
||||
.hx-music-dot-running {
|
||||
background-color: var(--color-accent);
|
||||
animation: hx-music-pulse 1.2s ease-in-out infinite;
|
||||
}
|
||||
.hx-music-dot-failed {
|
||||
background-color: var(--color-accent-hot);
|
||||
}
|
||||
.hx-music-swatch {
|
||||
flex-shrink: 0;
|
||||
width: 0.9rem;
|
||||
height: 0.9rem;
|
||||
border-radius: 0.2rem;
|
||||
}
|
||||
@keyframes hx-music-pulse {
|
||||
50% {
|
||||
opacity: 0.25;
|
||||
}
|
||||
}
|
||||
@media (prefers-reduced-motion: reduce) {
|
||||
.hx-music-dot-queued,
|
||||
.hx-music-dot-running {
|
||||
animation: none;
|
||||
}
|
||||
}
|
||||
|
||||
/* -- Documentation (/docs) ------------------------------------------------ */
|
||||
|
||||
@@ -43,7 +43,7 @@ export default function App() {
|
||||
<Route path="/chat" element={<Chat />} />
|
||||
<Route path="/images" element={<Images />} />
|
||||
<Route path="/music" element={<Music />} />
|
||||
<Route path="/playground" element={<Playground />} />
|
||||
<Route path="/decisions" element={<Playground />} />
|
||||
{/* Documentation. The splat route lets the sidebar tree be
|
||||
arbitrarily deep without enumerating routes here — the
|
||||
page resolves the slug against the content tree. */}
|
||||
|
||||
@@ -115,12 +115,12 @@ const Header: React.FC = () => {
|
||||
{t("nav.music")}
|
||||
</NavLink>
|
||||
<NavLink
|
||||
to="/playground"
|
||||
to="/decisions"
|
||||
className={({ isActive }): string =>
|
||||
isActive ? "nav-link active" : "nav-link"
|
||||
}
|
||||
>
|
||||
{t("nav.playground")}
|
||||
{t("nav.decisions")}
|
||||
</NavLink>
|
||||
{/* Documentation is reachable on every host — this only
|
||||
decides whether it is advertised. The pages are still
|
||||
|
||||
@@ -106,6 +106,20 @@ export interface GeneratedImage {
|
||||
createdAt: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Local organisation of music jobs (#402). The jobs and their audio live on
|
||||
* the server; which project a song is filed under, and any name given to
|
||||
* it, are this browser's own, like chat projects.
|
||||
*/
|
||||
export interface MusicTrack {
|
||||
/** The server's job id. */
|
||||
jobId: string;
|
||||
owner: string;
|
||||
projectId: string | null; // null → filed by date
|
||||
/** A name the user gave the song; the prompt is shown otherwise. */
|
||||
title?: string;
|
||||
}
|
||||
|
||||
/** Small key/value store: fingerprint, active conversation, anon usage. */
|
||||
export interface Meta {
|
||||
key: string;
|
||||
@@ -131,6 +145,8 @@ class HelexaDB extends Dexie {
|
||||
messages!: Table<Message, string>;
|
||||
meta!: Table<Meta, string>;
|
||||
images!: Table<GeneratedImage, string>;
|
||||
musicProjects!: Table<Project, string>;
|
||||
musicTracks!: Table<MusicTrack, string>;
|
||||
|
||||
constructor() {
|
||||
super("helexa");
|
||||
@@ -166,6 +182,13 @@ class HelexaDB extends Dexie {
|
||||
}
|
||||
}),
|
||||
);
|
||||
// v4 adds music projects and per-song filing (#402). Additive. Music
|
||||
// projects are their own store: a chat topic and a music project are
|
||||
// different things, and sharing one list would mix them.
|
||||
this.version(4).stores({
|
||||
musicProjects: "id, owner, [owner+archived], updatedAt",
|
||||
musicTracks: "jobId, owner, projectId",
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -9,6 +9,7 @@ import {
|
||||
type GeneratedImage,
|
||||
type Message,
|
||||
type MessageRole,
|
||||
type MusicTrack,
|
||||
type Project,
|
||||
type TurnBlock,
|
||||
} from "./db";
|
||||
@@ -193,3 +194,66 @@ export async function saveImage(
|
||||
export async function deleteImage(id: string): Promise<void> {
|
||||
await db.images.delete(id);
|
||||
}
|
||||
|
||||
// ── music projects (#402) ───────────────────────────────────────────
|
||||
|
||||
export async function listMusicProjects(owner: string): Promise<Project[]> {
|
||||
const rows = await db.musicProjects.where({ owner }).toArray();
|
||||
return rows
|
||||
.filter((p) => !p.archived)
|
||||
.sort((a, b) => a.sortOrder - b.sortOrder || a.createdAt - b.createdAt);
|
||||
}
|
||||
|
||||
export async function createMusicProject(owner: string, name: string): Promise<string> {
|
||||
const id = uuid();
|
||||
const ts = now();
|
||||
await db.musicProjects.add({
|
||||
id,
|
||||
owner,
|
||||
name,
|
||||
createdAt: ts,
|
||||
updatedAt: ts,
|
||||
archived: false,
|
||||
sortOrder: ts,
|
||||
});
|
||||
return id;
|
||||
}
|
||||
|
||||
export async function renameMusicProject(id: string, name: string): Promise<void> {
|
||||
await db.musicProjects.update(id, { name, updatedAt: now() });
|
||||
}
|
||||
|
||||
export async function archiveMusicProject(id: string): Promise<void> {
|
||||
// Its songs go back to the date sections rather than disappearing.
|
||||
await db.transaction("rw", db.musicProjects, db.musicTracks, async () => {
|
||||
await db.musicProjects.update(id, { archived: true, updatedAt: now() });
|
||||
await db.musicTracks.where({ projectId: id }).modify({ projectId: null });
|
||||
});
|
||||
}
|
||||
|
||||
export async function listMusicTracks(owner: string): Promise<MusicTrack[]> {
|
||||
return db.musicTracks.where({ owner }).toArray();
|
||||
}
|
||||
|
||||
/** File a song under a project, or back under its date with `null`. */
|
||||
export async function moveMusicTrack(
|
||||
owner: string,
|
||||
jobId: string,
|
||||
projectId: string | null,
|
||||
): Promise<void> {
|
||||
const existing = await db.musicTracks.get(jobId);
|
||||
await db.musicTracks.put({ ...existing, jobId, owner, projectId });
|
||||
}
|
||||
|
||||
export async function renameMusicTrack(owner: string, jobId: string, title: string): Promise<void> {
|
||||
const existing = await db.musicTracks.get(jobId);
|
||||
await db.musicTracks.put({ projectId: null, ...existing, jobId, owner, title });
|
||||
}
|
||||
|
||||
/** Drop filing for songs the server no longer has (expired or never seen). */
|
||||
export async function pruneMusicTracks(owner: string, liveJobIds: Set<string>): Promise<void> {
|
||||
await db.musicTracks
|
||||
.where({ owner })
|
||||
.filter((t) => !liveJobIds.has(t.jobId))
|
||||
.delete();
|
||||
}
|
||||
|
||||
@@ -64,9 +64,9 @@ describe("decisions translations", () => {
|
||||
}
|
||||
});
|
||||
|
||||
it.each(Object.keys(COMMON))("%s names the playground in the nav", (lang) => {
|
||||
it.each(Object.keys(COMMON))("%s names the decisions page in the nav", (lang) => {
|
||||
const nav = COMMON[lang].nav as Tree;
|
||||
expect(typeof nav.playground).toBe("string");
|
||||
expect((nav.playground as string).trim()).not.toBe("");
|
||||
expect(typeof nav.decisions).toBe("string");
|
||||
expect((nav.decisions as string).trim()).not.toBe("");
|
||||
});
|
||||
});
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
"chat": "ውይይት",
|
||||
"images": "ምስሎች",
|
||||
"music": "ሙዚቃ",
|
||||
"playground": "የሙከራ ሜዳ",
|
||||
"decisions": "ውሳኔዎች",
|
||||
"mission": "ተልእኮ",
|
||||
"login": "ይግቡ",
|
||||
"register": "ይመዝገቡ",
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"page": {
|
||||
"title": "የሙከራ ሜዳ",
|
||||
"lead": "ለውሳኔ ሞዴል ስለማንኛውም ጽሑፍ ወይም JSON ዓይነት ያላቸው ጥያቄዎችን ይጠይቁ፤ በአንድ ዙር የተስተካከሉ መልሶችን ያግኙ — ያለ የመነጨ ጽሑፍ።",
|
||||
"title": "ውሳኔዎች",
|
||||
"lead": "ስለማንኛውም ጽሑፍ ወይም JSON አዎ/አይደለም፣ ባለብዙ ምርጫ ወይም የነጥብ ጥያቄዎችን ይጠይቁ፤ በተስተካከለ እርግጠኝነት ፈጣን መልሶችን ያግኙ። እንደዚህ ያሉ ሞዴሎች አንዳንድ ጊዜ System One ሞዴሎች ይባላሉ።",
|
||||
"docsLink": "እንዴት እንደሚሠራ",
|
||||
"side": "ውጤቶች፣ ምሳሌዎች እና ኮድ",
|
||||
"tabs": {
|
||||
|
||||
@@ -33,5 +33,21 @@
|
||||
"cancel": "ሰርዝ",
|
||||
"cost": "ወጪ",
|
||||
"errorGeneric": "ዘፈኑን መፍጠር አልተቻለም።",
|
||||
"errorLoad": "ድምጹን መጫን አልተቻለም።"
|
||||
"errorLoad": "ድምጹን መጫን አልተቻለም።",
|
||||
"history": "ታሪክ",
|
||||
"period": {
|
||||
"today": "ዛሬ",
|
||||
"yesterday": "ትናንት",
|
||||
"week": "ያለፉት 7 ቀናት",
|
||||
"month": "ያለፉት 30 ቀናት",
|
||||
"older": "የቆዩ"
|
||||
},
|
||||
"length": "ርዝመት",
|
||||
"tags": "የስልት ሃሳቦች",
|
||||
"submitting": "ዘፈንዎ ወደ ስቱዲዮ እየተላከ ነው…",
|
||||
"needStyle": "ዘፈን ለመፍጠር ስልት ይግለጹ።",
|
||||
"needLength": "ከ{{min}} እስከ {{max}} ሰከንድ ያለ ርዝመት ይምረጡ።",
|
||||
"close": "ዝጋ",
|
||||
"pause": "ለአፍታ አቁም",
|
||||
"seek": "በዘፈኑ ውስጥ ያለ ቦታ"
|
||||
}
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
"chat": "المحادثة",
|
||||
"images": "الصور",
|
||||
"music": "الموسيقى",
|
||||
"playground": "ساحة التجربة",
|
||||
"decisions": "القرارات",
|
||||
"mission": "المهمة",
|
||||
"login": "تسجيل الدخول",
|
||||
"register": "إنشاء حساب",
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"page": {
|
||||
"title": "ساحة التجربة",
|
||||
"lead": "اطرح على نموذج قرارات أسئلة محدّدة النوع عن أي نص أو JSON، واحصل على إجابات معايَرة في تمريرة واحدة — دون نص مولَّد.",
|
||||
"title": "القرارات",
|
||||
"lead": "اطرح أسئلة بنعم/لا أو اختيار من متعدد أو تقييم حول أي نص أو JSON، واحصل فورًا على إجابات بثقة معايَرة. تُسمّى هذه النماذج أحيانًا نماذج System One.",
|
||||
"docsLink": "كيف يعمل",
|
||||
"side": "النتائج والأمثلة والشيفرة",
|
||||
"tabs": {
|
||||
|
||||
@@ -33,5 +33,21 @@
|
||||
"cancel": "إلغاء",
|
||||
"cost": "التكلفة",
|
||||
"errorGeneric": "تعذّر إنشاء الأغنية.",
|
||||
"errorLoad": "تعذّر تحميل الصوت."
|
||||
"errorLoad": "تعذّر تحميل الصوت.",
|
||||
"history": "السجل",
|
||||
"period": {
|
||||
"today": "اليوم",
|
||||
"yesterday": "أمس",
|
||||
"week": "آخر 7 أيام",
|
||||
"month": "آخر 30 يومًا",
|
||||
"older": "أقدم"
|
||||
},
|
||||
"length": "المدة",
|
||||
"tags": "أفكار للأسلوب",
|
||||
"submitting": "جارٍ إرسال أغنيتك إلى الاستوديو…",
|
||||
"needStyle": "صِف أسلوبًا لإنشاء أغنية.",
|
||||
"needLength": "اختر مدة بين {{min}} و{{max}} ثانية.",
|
||||
"close": "إغلاق",
|
||||
"pause": "إيقاف مؤقت",
|
||||
"seek": "الموضع في الأغنية"
|
||||
}
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
"chat": "Чат",
|
||||
"images": "Изображения",
|
||||
"music": "Музика",
|
||||
"playground": "Пясъчник",
|
||||
"decisions": "Решения",
|
||||
"mission": "Мисия",
|
||||
"login": "Вход",
|
||||
"register": "Регистрация",
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"page": {
|
||||
"title": "Пясъчник",
|
||||
"lead": "Задавайте на модел за вземане на решения типизирани въпроси за всякакъв текст или JSON и получавайте калибрирани отговори с едно преминаване — без генериран текст.",
|
||||
"title": "Решения",
|
||||
"lead": "Задавайте въпроси с да/не, с избор между възможности или с оценка за всякакъв текст или JSON и получавайте мигновено отговори с калибрирана увереност. Такива модели понякога се наричат модели System One.",
|
||||
"docsLink": "Как работи",
|
||||
"side": "Резултати, примери и код",
|
||||
"tabs": {
|
||||
|
||||
@@ -33,5 +33,21 @@
|
||||
"cancel": "Отказ",
|
||||
"cost": "Цена",
|
||||
"errorGeneric": "Песента не можа да бъде създадена.",
|
||||
"errorLoad": "Аудиото не можа да бъде заредено."
|
||||
"errorLoad": "Аудиото не можа да бъде заредено.",
|
||||
"history": "История",
|
||||
"period": {
|
||||
"today": "Днес",
|
||||
"yesterday": "Вчера",
|
||||
"week": "Последните 7 дни",
|
||||
"month": "Последните 30 дни",
|
||||
"older": "По-стари"
|
||||
},
|
||||
"length": "Продължителност",
|
||||
"tags": "Идеи за стил",
|
||||
"submitting": "Изпращане на песента ви към студиото…",
|
||||
"needStyle": "Опишете стил, за да генерирате песен.",
|
||||
"needLength": "Изберете продължителност между {{min}} и {{max}} секунди.",
|
||||
"close": "Затвори",
|
||||
"pause": "Пауза",
|
||||
"seek": "Позиция в песента"
|
||||
}
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
"chat": "Xat",
|
||||
"images": "Imatges",
|
||||
"music": "Música",
|
||||
"playground": "Espai de proves",
|
||||
"decisions": "Decisions",
|
||||
"mission": "Missió",
|
||||
"login": "Inicia la sessió",
|
||||
"register": "Registra't",
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"page": {
|
||||
"title": "Espai de proves",
|
||||
"lead": "Feu a un model de decisió preguntes tipades sobre qualsevol text o JSON i obtingueu respostes calibrades en una sola passada, sense text generat.",
|
||||
"title": "Decisions",
|
||||
"lead": "Feu preguntes de sí/no, d’opció múltiple o de puntuació sobre qualsevol text o JSON i obteniu a l’instant respostes amb una confiança calibrada. De vegades se’n diu models System One.",
|
||||
"docsLink": "Com funciona",
|
||||
"side": "Resultats, exemples i codi",
|
||||
"tabs": {
|
||||
|
||||
@@ -33,5 +33,21 @@
|
||||
"cancel": "Cancel·la",
|
||||
"cost": "Cost",
|
||||
"errorGeneric": "No s’ha pogut generar la cançó.",
|
||||
"errorLoad": "No s’ha pogut carregar l’àudio."
|
||||
"errorLoad": "No s’ha pogut carregar l’àudio.",
|
||||
"history": "Historial",
|
||||
"period": {
|
||||
"today": "Avui",
|
||||
"yesterday": "Ahir",
|
||||
"week": "Darrers 7 dies",
|
||||
"month": "Darrers 30 dies",
|
||||
"older": "Més antics"
|
||||
},
|
||||
"length": "Durada",
|
||||
"tags": "Idees d’estil",
|
||||
"submitting": "S’està enviant la cançó a l’estudi…",
|
||||
"needStyle": "Descriviu un estil per generar una cançó.",
|
||||
"needLength": "Trieu una durada entre {{min}} i {{max}} segons.",
|
||||
"close": "Tanca",
|
||||
"pause": "Pausa",
|
||||
"seek": "Posició a la cançó"
|
||||
}
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
"chat": "Chat",
|
||||
"images": "Obrázky",
|
||||
"music": "Hudba",
|
||||
"playground": "Pískoviště",
|
||||
"decisions": "Rozhodnutí",
|
||||
"mission": "Mise",
|
||||
"login": "Přihlásit se",
|
||||
"register": "Registrovat se",
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"page": {
|
||||
"title": "Pískoviště",
|
||||
"lead": "Pokládejte rozhodovacímu modelu typované otázky o libovolném textu nebo JSON a dostávejte kalibrované odpovědi v jediném průchodu — bez generovaného textu.",
|
||||
"title": "Rozhodnutí",
|
||||
"lead": "Ptejte se otázkami ano/ne, s výběrem z možností nebo s hodnocením na libovolný text či JSON a okamžitě získejte odpovědi s kalibrovanou jistotou. Takovým modelům se někdy říká modely System One.",
|
||||
"docsLink": "Jak to funguje",
|
||||
"side": "Výsledky, příklady a kód",
|
||||
"tabs": {
|
||||
|
||||
@@ -33,5 +33,21 @@
|
||||
"cancel": "Zrušit",
|
||||
"cost": "Cena",
|
||||
"errorGeneric": "Skladbu se nepodařilo vygenerovat.",
|
||||
"errorLoad": "Zvuk se nepodařilo načíst."
|
||||
"errorLoad": "Zvuk se nepodařilo načíst.",
|
||||
"history": "Historie",
|
||||
"period": {
|
||||
"today": "Dnes",
|
||||
"yesterday": "Včera",
|
||||
"week": "Posledních 7 dní",
|
||||
"month": "Posledních 30 dní",
|
||||
"older": "Starší"
|
||||
},
|
||||
"length": "Délka",
|
||||
"tags": "Nápady na styl",
|
||||
"submitting": "Odesílám skladbu do studia…",
|
||||
"needStyle": "Popište styl, abyste mohli vygenerovat skladbu.",
|
||||
"needLength": "Zvolte délku mezi {{min}} a {{max}} sekundami.",
|
||||
"close": "Zavřít",
|
||||
"pause": "Pozastavit",
|
||||
"seek": "Pozice ve skladbě"
|
||||
}
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
"chat": "Chat",
|
||||
"images": "Billeder",
|
||||
"music": "Musik",
|
||||
"playground": "Legeplads",
|
||||
"decisions": "Beslutninger",
|
||||
"mission": "Mission",
|
||||
"login": "Log ind",
|
||||
"register": "Opret konto",
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"page": {
|
||||
"title": "Legeplads",
|
||||
"lead": "Stil en beslutningsmodel typede spørgsmål om enhver tekst eller JSON, og få kalibrerede svar i ét gennemløb — uden genereret tekst.",
|
||||
"title": "Beslutninger",
|
||||
"lead": "Stil ja/nej-, multiple choice- eller pointspørgsmål om enhver tekst eller JSON, og få straks svar med kalibreret sikkerhed. Den slags modeller kaldes nogle gange System One-modeller.",
|
||||
"docsLink": "Sådan virker det",
|
||||
"side": "Resultater, eksempler og kode",
|
||||
"tabs": {
|
||||
|
||||
@@ -33,5 +33,21 @@
|
||||
"cancel": "Annullér",
|
||||
"cost": "Pris",
|
||||
"errorGeneric": "Sangen kunne ikke genereres.",
|
||||
"errorLoad": "Lyden kunne ikke indlæses."
|
||||
"errorLoad": "Lyden kunne ikke indlæses.",
|
||||
"history": "Historik",
|
||||
"period": {
|
||||
"today": "I dag",
|
||||
"yesterday": "I går",
|
||||
"week": "Seneste 7 dage",
|
||||
"month": "Seneste 30 dage",
|
||||
"older": "Ældre"
|
||||
},
|
||||
"length": "Længde",
|
||||
"tags": "Stilidéer",
|
||||
"submitting": "Sender din sang til studiet…",
|
||||
"needStyle": "Beskriv en stil for at generere en sang.",
|
||||
"needLength": "Vælg en længde mellem {{min}} og {{max}} sekunder.",
|
||||
"close": "Luk",
|
||||
"pause": "Pause",
|
||||
"seek": "Position i sangen"
|
||||
}
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
"chat": "Chat",
|
||||
"images": "Bilder",
|
||||
"music": "Musik",
|
||||
"playground": "Playground",
|
||||
"decisions": "Entscheidungen",
|
||||
"mission": "Mission",
|
||||
"login": "Anmelden",
|
||||
"register": "Registrieren",
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"page": {
|
||||
"title": "Playground",
|
||||
"lead": "Stellen Sie einem Entscheidungsmodell typisierte Fragen zu beliebigem Text oder JSON und erhalten Sie kalibrierte Antworten in einem Durchlauf – ganz ohne generierten Text.",
|
||||
"title": "Entscheidungen",
|
||||
"lead": "Stellen Sie Ja/Nein-, Auswahl- oder Bewertungsfragen zu beliebigem Text oder JSON und erhalten Sie sofort Antworten mit kalibrierter Konfidenz. Solche Modelle werden manchmal System-One-Modelle genannt.",
|
||||
"docsLink": "So funktioniert es",
|
||||
"side": "Ergebnisse, Beispiele und Code",
|
||||
"tabs": {
|
||||
|
||||
@@ -33,5 +33,21 @@
|
||||
"cancel": "Abbrechen",
|
||||
"cost": "Kosten",
|
||||
"errorGeneric": "Der Song konnte nicht erzeugt werden.",
|
||||
"errorLoad": "Die Audiodatei konnte nicht geladen werden."
|
||||
"errorLoad": "Die Audiodatei konnte nicht geladen werden.",
|
||||
"history": "Verlauf",
|
||||
"period": {
|
||||
"today": "Heute",
|
||||
"yesterday": "Gestern",
|
||||
"week": "Letzte 7 Tage",
|
||||
"month": "Letzte 30 Tage",
|
||||
"older": "Älter"
|
||||
},
|
||||
"length": "Länge",
|
||||
"tags": "Stilideen",
|
||||
"submitting": "Ihr Song wird ans Studio gesendet…",
|
||||
"needStyle": "Beschreiben Sie einen Stil, um einen Song zu erzeugen.",
|
||||
"needLength": "Wählen Sie eine Länge zwischen {{min}} und {{max}} Sekunden.",
|
||||
"close": "Schließen",
|
||||
"pause": "Pause",
|
||||
"seek": "Position im Song"
|
||||
}
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
"chat": "Συνομιλία",
|
||||
"images": "Εικόνες",
|
||||
"music": "Μουσική",
|
||||
"playground": "Πεδίο δοκιμών",
|
||||
"decisions": "Αποφάσεις",
|
||||
"mission": "Αποστολή",
|
||||
"login": "Σύνδεση",
|
||||
"register": "Εγγραφή",
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"page": {
|
||||
"title": "Πεδίο δοκιμών",
|
||||
"lead": "Κάντε σε ένα μοντέλο αποφάσεων τυποποιημένες ερωτήσεις για οποιοδήποτε κείμενο ή JSON και λάβετε βαθμονομημένες απαντήσεις σε ένα μόνο πέρασμα — χωρίς παραγόμενο κείμενο.",
|
||||
"title": "Αποφάσεις",
|
||||
"lead": "Κάντε ερωτήσεις ναι/όχι, πολλαπλής επιλογής ή βαθμολόγησης για οποιοδήποτε κείμενο ή JSON και λάβετε αμέσως απαντήσεις με βαθμονομημένη βεβαιότητα. Τέτοια μοντέλα ονομάζονται μερικές φορές μοντέλα System One.",
|
||||
"docsLink": "Πώς λειτουργεί",
|
||||
"side": "Αποτελέσματα, παραδείγματα και κώδικας",
|
||||
"tabs": {
|
||||
|
||||
@@ -33,5 +33,21 @@
|
||||
"cancel": "Ακύρωση",
|
||||
"cost": "Κόστος",
|
||||
"errorGeneric": "Δεν ήταν δυνατή η δημιουργία του τραγουδιού.",
|
||||
"errorLoad": "Δεν ήταν δυνατή η φόρτωση του ήχου."
|
||||
"errorLoad": "Δεν ήταν δυνατή η φόρτωση του ήχου.",
|
||||
"history": "Ιστορικό",
|
||||
"period": {
|
||||
"today": "Σήμερα",
|
||||
"yesterday": "Χθες",
|
||||
"week": "Τελευταίες 7 ημέρες",
|
||||
"month": "Τελευταίες 30 ημέρες",
|
||||
"older": "Παλαιότερα"
|
||||
},
|
||||
"length": "Διάρκεια",
|
||||
"tags": "Ιδέες ύφους",
|
||||
"submitting": "Αποστολή του τραγουδιού σας στο στούντιο…",
|
||||
"needStyle": "Περιγράψτε ένα ύφος για να δημιουργήσετε τραγούδι.",
|
||||
"needLength": "Επιλέξτε διάρκεια από {{min}} έως {{max}} δευτερόλεπτα.",
|
||||
"close": "Κλείσιμο",
|
||||
"pause": "Παύση",
|
||||
"seek": "Θέση στο τραγούδι"
|
||||
}
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
"chat": "Chat",
|
||||
"images": "Images",
|
||||
"music": "Music",
|
||||
"playground": "Playground",
|
||||
"decisions": "Decisions",
|
||||
"mission": "Mission",
|
||||
"login": "Sign in",
|
||||
"register": "Sign up",
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"page": {
|
||||
"title": "Playground",
|
||||
"lead": "Ask a decision model typed questions about any text or JSON, and get calibrated answers in one pass — no generated text.",
|
||||
"title": "Decisions",
|
||||
"lead": "Ask yes/no, multiple-choice or scoring questions about any text or JSON, and get instant answers with calibrated confidence. Models like this are sometimes called System One models.",
|
||||
"docsLink": "How it works",
|
||||
"side": "Results, examples and code",
|
||||
"tabs": {
|
||||
|
||||
@@ -33,5 +33,21 @@
|
||||
"cancel": "Cancel",
|
||||
"cost": "Cost",
|
||||
"errorGeneric": "The song could not be generated.",
|
||||
"errorLoad": "The audio could not be loaded."
|
||||
"errorLoad": "The audio could not be loaded.",
|
||||
"history": "History",
|
||||
"period": {
|
||||
"today": "Today",
|
||||
"yesterday": "Yesterday",
|
||||
"week": "Previous 7 days",
|
||||
"month": "Previous 30 days",
|
||||
"older": "Older"
|
||||
},
|
||||
"length": "Length",
|
||||
"tags": "Style ideas",
|
||||
"submitting": "Sending your song to the studio…",
|
||||
"needStyle": "Describe a style to generate a song.",
|
||||
"needLength": "Choose a length between {{min}} and {{max}} seconds.",
|
||||
"close": "Close",
|
||||
"pause": "Pause",
|
||||
"seek": "Position in the song"
|
||||
}
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
"chat": "Chat",
|
||||
"images": "Imágenes",
|
||||
"music": "Música",
|
||||
"playground": "Zona de pruebas",
|
||||
"decisions": "Decisiones",
|
||||
"mission": "Misión",
|
||||
"login": "Iniciar sesión",
|
||||
"register": "Registrarse",
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"page": {
|
||||
"title": "Zona de pruebas",
|
||||
"lead": "Haga a un modelo de decisión preguntas tipadas sobre cualquier texto o JSON y obtenga respuestas calibradas en una sola pasada, sin texto generado.",
|
||||
"title": "Decisiones",
|
||||
"lead": "Haga preguntas de sí/no, de opción múltiple o de puntuación sobre cualquier texto o JSON y obtenga al instante respuestas con una confianza calibrada. A estos modelos a veces se les llama modelos System One.",
|
||||
"docsLink": "Cómo funciona",
|
||||
"side": "Resultados, ejemplos y código",
|
||||
"tabs": {
|
||||
|
||||
@@ -33,5 +33,21 @@
|
||||
"cancel": "Cancelar",
|
||||
"cost": "Coste",
|
||||
"errorGeneric": "No se pudo generar la canción.",
|
||||
"errorLoad": "No se pudo cargar el audio."
|
||||
"errorLoad": "No se pudo cargar el audio.",
|
||||
"history": "Historial",
|
||||
"period": {
|
||||
"today": "Hoy",
|
||||
"yesterday": "Ayer",
|
||||
"week": "Últimos 7 días",
|
||||
"month": "Últimos 30 días",
|
||||
"older": "Anteriores"
|
||||
},
|
||||
"length": "Duración",
|
||||
"tags": "Ideas de estilo",
|
||||
"submitting": "Enviando su canción al estudio…",
|
||||
"needStyle": "Describa un estilo para generar una canción.",
|
||||
"needLength": "Elija una duración entre {{min}} y {{max}} segundos.",
|
||||
"close": "Cerrar",
|
||||
"pause": "Pausa",
|
||||
"seek": "Posición en la canción"
|
||||
}
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
"chat": "Vestlus",
|
||||
"images": "Pildid",
|
||||
"music": "Muusika",
|
||||
"playground": "Mänguväljak",
|
||||
"decisions": "Otsused",
|
||||
"mission": "Missioon",
|
||||
"login": "Logi sisse",
|
||||
"register": "Registreeru",
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"page": {
|
||||
"title": "Mänguväljak",
|
||||
"lead": "Esitage otsustusmudelile tüübitud küsimusi mis tahes teksti või JSON-i kohta ja saage kalibreeritud vastused ühe läbimisega — ilma genereeritud tekstita.",
|
||||
"title": "Otsused",
|
||||
"lead": "Esita jah/ei-, valikvastustega või hindamisküsimusi mis tahes teksti või JSON-i kohta ja saa kohe vastused kalibreeritud kindlusega. Selliseid mudeleid nimetatakse mõnikord System One mudeliteks.",
|
||||
"docsLink": "Kuidas see töötab",
|
||||
"side": "Tulemused, näited ja kood",
|
||||
"tabs": {
|
||||
|
||||
@@ -33,5 +33,21 @@
|
||||
"cancel": "Tühista",
|
||||
"cost": "Kulu",
|
||||
"errorGeneric": "Laulu ei õnnestunud luua.",
|
||||
"errorLoad": "Heli ei õnnestunud laadida."
|
||||
"errorLoad": "Heli ei õnnestunud laadida.",
|
||||
"history": "Ajalugu",
|
||||
"period": {
|
||||
"today": "Täna",
|
||||
"yesterday": "Eile",
|
||||
"week": "Viimased 7 päeva",
|
||||
"month": "Viimased 30 päeva",
|
||||
"older": "Vanemad"
|
||||
},
|
||||
"length": "Pikkus",
|
||||
"tags": "Stiiliideed",
|
||||
"submitting": "Saadan su laulu stuudiosse…",
|
||||
"needStyle": "Laulu loomiseks kirjelda stiili.",
|
||||
"needLength": "Vali pikkus vahemikus {{min}}–{{max}} sekundit.",
|
||||
"close": "Sulge",
|
||||
"pause": "Paus",
|
||||
"seek": "Asukoht laulus"
|
||||
}
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
"chat": "گفتگو",
|
||||
"images": "تصاویر",
|
||||
"music": "موسیقی",
|
||||
"playground": "زمین آزمایش",
|
||||
"decisions": "تصمیمها",
|
||||
"mission": "مأموریت",
|
||||
"login": "ورود",
|
||||
"register": "ثبتنام",
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"page": {
|
||||
"title": "زمین آزمایش",
|
||||
"lead": "از یک مدل تصمیمگیری دربارهٔ هر متن یا JSON پرسشهای نوعدار بپرسید و در یک گذر پاسخهای کالیبرهشده بگیرید — بدون متن تولیدشده.",
|
||||
"title": "تصمیمها",
|
||||
"lead": "دربارهٔ هر متن یا JSON پرسشهای بله/خیر، چندگزینهای یا امتیازدهی بپرسید و فوراً پاسخهایی با اطمینان کالیبرهشده بگیرید. به چنین مدلهایی گاهی مدلهای System One میگویند.",
|
||||
"docsLink": "چطور کار میکند",
|
||||
"side": "نتایج، نمونهها و کد",
|
||||
"tabs": {
|
||||
|
||||
@@ -33,5 +33,21 @@
|
||||
"cancel": "لغو",
|
||||
"cost": "هزینه",
|
||||
"errorGeneric": "ساخت آهنگ ممکن نشد.",
|
||||
"errorLoad": "بارگذاری صدا ممکن نشد."
|
||||
"errorLoad": "بارگذاری صدا ممکن نشد.",
|
||||
"history": "تاریخچه",
|
||||
"period": {
|
||||
"today": "امروز",
|
||||
"yesterday": "دیروز",
|
||||
"week": "۷ روز گذشته",
|
||||
"month": "۳۰ روز گذشته",
|
||||
"older": "قدیمیتر"
|
||||
},
|
||||
"length": "مدت",
|
||||
"tags": "ایدههای سبک",
|
||||
"submitting": "در حال ارسال آهنگ شما به استودیو…",
|
||||
"needStyle": "برای ساخت آهنگ، یک سبک توصیف کنید.",
|
||||
"needLength": "مدتی بین {{min}} تا {{max}} ثانیه انتخاب کنید.",
|
||||
"close": "بستن",
|
||||
"pause": "مکث",
|
||||
"seek": "موقعیت در آهنگ"
|
||||
}
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
"chat": "Keskustelu",
|
||||
"images": "Kuvat",
|
||||
"music": "Musiikki",
|
||||
"playground": "Hiekkalaatikko",
|
||||
"decisions": "Päätökset",
|
||||
"mission": "Missio",
|
||||
"login": "Kirjaudu sisään",
|
||||
"register": "Rekisteröidy",
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"page": {
|
||||
"title": "Hiekkalaatikko",
|
||||
"lead": "Esitä päätösmallille tyypitettyjä kysymyksiä mistä tahansa tekstistä tai JSONista ja saat kalibroidut vastaukset yhdellä läpikäynnillä — ilman tuotettua tekstiä.",
|
||||
"title": "Päätökset",
|
||||
"lead": "Esitä kyllä/ei-, monivalinta- tai pisteytyskysymyksiä mistä tahansa tekstistä tai JSONista ja saat heti vastaukset kalibroidulla varmuudella. Tällaisia malleja kutsutaan joskus System One -malleiksi.",
|
||||
"docsLink": "Miten se toimii",
|
||||
"side": "Tulokset, esimerkit ja koodi",
|
||||
"tabs": {
|
||||
|
||||
@@ -33,5 +33,21 @@
|
||||
"cancel": "Peruuta",
|
||||
"cost": "Hinta",
|
||||
"errorGeneric": "Kappaletta ei voitu luoda.",
|
||||
"errorLoad": "Ääntä ei voitu ladata."
|
||||
"errorLoad": "Ääntä ei voitu ladata.",
|
||||
"history": "Historia",
|
||||
"period": {
|
||||
"today": "Tänään",
|
||||
"yesterday": "Eilen",
|
||||
"week": "Edelliset 7 päivää",
|
||||
"month": "Edelliset 30 päivää",
|
||||
"older": "Vanhemmat"
|
||||
},
|
||||
"length": "Pituus",
|
||||
"tags": "Tyyli-ideoita",
|
||||
"submitting": "Kappaletta lähetetään studioon…",
|
||||
"needStyle": "Kuvaile tyyli luodaksesi kappaleen.",
|
||||
"needLength": "Valitse pituus väliltä {{min}}–{{max}} sekuntia.",
|
||||
"close": "Sulje",
|
||||
"pause": "Tauko",
|
||||
"seek": "Kohta kappaleessa"
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user