mirror of
https://github.com/mudler/LocalAI.git
synced 2026-05-20 14:46:38 -04:00
Compare commits
143 Commits
feat/vllm-
...
feat/buun-
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9787bee48b | ||
|
|
42754d33b9 | ||
|
|
7f2b7e4ace | ||
|
|
6233feb190 | ||
|
|
d6bf3a4969 | ||
|
|
b27d38a53d | ||
|
|
45756b19dc | ||
|
|
cd6079b2f3 | ||
|
|
3db60b57e6 | ||
|
|
13734ae9fa | ||
|
|
c0920f3273 | ||
|
|
7c1934b183 | ||
|
|
5e062b4d1f | ||
|
|
4906cbad04 | ||
|
|
c755cd5ab5 | ||
|
|
0fb04f7ac3 | ||
|
|
d9d7b5c29b | ||
|
|
f877942d97 | ||
|
|
f5eb13d3c2 | ||
|
|
c1f923b2bc | ||
|
|
ed648b3b4e | ||
|
|
3ce5248126 | ||
|
|
04f1a0285d | ||
|
|
181ebb6df4 | ||
|
|
1c59165d63 | ||
|
|
eb00d9b178 | ||
|
|
2068b6f43c | ||
|
|
eb01c77214 | ||
|
|
bb4fda6f0e | ||
|
|
f0c92610a1 | ||
|
|
bbeacf140d | ||
|
|
6820ec468f | ||
|
|
20baec77ab | ||
|
|
d16f19f1eb | ||
|
|
cd7b035716 | ||
|
|
0f3bb2d647 | ||
|
|
607efe5a4c | ||
|
|
7d8c1d5e45 | ||
|
|
d18d434bb2 | ||
|
|
39573ecd2a | ||
|
|
a7dbb2a83d | ||
|
|
3ad9b16c29 | ||
|
|
c806d5ab73 | ||
|
|
47efaf5b43 | ||
|
|
315b634a91 | ||
|
|
6b245299d7 | ||
|
|
677c0315c1 | ||
|
|
478522ce4d | ||
|
|
c54897ad44 | ||
|
|
8bb1e8f21f | ||
|
|
cd94a0b61a | ||
|
|
047bc48fa9 | ||
|
|
01bd8ae5d0 | ||
|
|
d9808769be | ||
|
|
5973c0a9df | ||
|
|
486b5e25a3 | ||
|
|
c66c41e8d7 | ||
|
|
02bb715c0a | ||
|
|
8ab56e2ad3 | ||
|
|
ecf85fde9e | ||
|
|
6480715a16 | ||
|
|
f683231811 | ||
|
|
960757f0e8 | ||
|
|
865fd552f5 | ||
|
|
cb77a5a4b9 | ||
|
|
60633c4dd5 | ||
|
|
9e44944cc1 | ||
|
|
372eb08dcf | ||
|
|
28091d626e | ||
|
|
cae79d9107 | ||
|
|
babbbc6ec8 | ||
|
|
3804497186 | ||
|
|
fda1c553a1 | ||
|
|
b27de08fff | ||
|
|
510f791ccc | ||
|
|
369c50a41c | ||
|
|
75a63f87d8 | ||
|
|
9cd8d7951f | ||
|
|
884bfb84c9 | ||
|
|
e94a9a8f10 | ||
|
|
054c4b4b45 | ||
|
|
6e49dba27c | ||
|
|
e463820566 | ||
|
|
8839a71c87 | ||
|
|
117f6430b8 | ||
|
|
7809c5f5d0 | ||
|
|
ad742738cb | ||
|
|
86c673fd94 | ||
|
|
c49feb546f | ||
|
|
844b0b760b | ||
|
|
55c05211d3 | ||
|
|
a90a8cf1d0 | ||
|
|
12b069f9bd | ||
|
|
48e87db400 | ||
|
|
7dbd9c056a | ||
|
|
7c5d6162f7 | ||
|
|
5837b14888 | ||
|
|
b6a68e5df4 | ||
|
|
c6dfb4acaf | ||
|
|
ec5935421c | ||
|
|
a0cbc46be9 | ||
|
|
b4e30692a2 | ||
|
|
61d34ccb11 | ||
|
|
7f88a3ba30 | ||
|
|
c4f309388e | ||
|
|
ab326a9c61 | ||
|
|
df2d25cee5 | ||
|
|
96cd561d9d | ||
|
|
08445b1b89 | ||
|
|
ad3c8c4832 | ||
|
|
6f0051301b | ||
|
|
8487058673 | ||
|
|
62862ca06b | ||
|
|
07e244d869 | ||
|
|
95efb8a562 | ||
|
|
410d100cc3 | ||
|
|
833b7e8557 | ||
|
|
87e6de1989 | ||
|
|
b361d2ddd6 | ||
|
|
1e4c4577bb | ||
|
|
98fd9d5cc6 | ||
|
|
0c725f5702 | ||
|
|
7661a4ffa5 | ||
|
|
24ad6e4be1 | ||
|
|
c0648b8836 | ||
|
|
a05c7def59 | ||
|
|
906acba8db | ||
|
|
4226ca4aee | ||
|
|
c6d5dc3374 | ||
|
|
7ce675af21 | ||
|
|
be1b8d56c9 | ||
|
|
97f087ed31 | ||
|
|
8691bbe663 | ||
|
|
7998f96f11 | ||
|
|
cada97ee46 | ||
|
|
3375ea1a2c | ||
|
|
0e7c0adee4 | ||
|
|
016da02845 | ||
|
|
daa0272f2e | ||
|
|
d67623230f | ||
|
|
0f90d17aac | ||
|
|
ea32b8953f | ||
|
|
bc7578bdb1 |
@@ -8,6 +8,7 @@ Create the backend directory under the appropriate location:
|
||||
- **Python backends**: `backend/python/<backend-name>/`
|
||||
- **Go backends**: `backend/go/<backend-name>/`
|
||||
- **C++ backends**: `backend/cpp/<backend-name>/`
|
||||
- **Rust backends**: `backend/rust/<backend-name>/`
|
||||
|
||||
For Python backends, you'll typically need:
|
||||
- `backend.py` - Main gRPC server implementation
|
||||
@@ -18,9 +19,22 @@ For Python backends, you'll typically need:
|
||||
- `run.sh` - Runtime script
|
||||
- `test.py` / `test.sh` - Test files
|
||||
|
||||
For Rust backends, you'll typically need (see `backend/rust/kokoros/` as a reference):
|
||||
- `Cargo.toml` - Crate manifest; depend on the upstream project as a submodule under `sources/`
|
||||
- `build.rs` - Invokes `tonic_build` to generate gRPC stubs from `backend/backend.proto` (use the `BACKEND_PROTO_PATH` env var so the Makefile can inject the canonical copy)
|
||||
- `src/` - The gRPC server implementation (implement `Backend` via `tonic`)
|
||||
- `Makefile` - Copies `backend.proto` into the crate, runs `cargo build --release`, then `package.sh`
|
||||
- `package.sh` - Uses `ldd` to bundle the binary's dynamic deps and `ld.so` into `package/lib/`
|
||||
- `run.sh` - Sets `LD_LIBRARY_PATH`/`SSL_CERT_DIR` and execs the binary via the bundled `lib/ld.so`
|
||||
- `sources/<UpstreamProject>/` - Git submodule with the upstream Rust crate
|
||||
|
||||
## 2. Add Build Configurations to `.github/workflows/backend.yml`
|
||||
|
||||
Add build matrix entries for each platform/GPU type you want to support. Look at similar backends (e.g., `chatterbox`, `faster-whisper`) for reference.
|
||||
Add build matrix entries for each platform/GPU type you want to support. Look at similar backends for reference — `chatterbox`/`faster-whisper` for Python, `piper`/`silero-vad` for Go, `kokoros` for Rust.
|
||||
|
||||
**Without an entry here no image is ever built or pushed, and the gallery entry in `backend/index.yaml` will point at a tag that does not exist.** The `dockerfile:` field must point at `./backend/Dockerfile.<lang>` matching the language bucket from step 1 (e.g. `Dockerfile.python`, `Dockerfile.golang`, `Dockerfile.rust`). The `tag-suffix` must match the `uri:` in the corresponding `backend/index.yaml` image entry exactly.
|
||||
|
||||
If you add a new language bucket, `scripts/changed-backends.js` also needs a branch in `inferBackendPath` so PR change-detection routes file edits correctly.
|
||||
|
||||
**Placement in file:**
|
||||
- CPU builds: Add after other CPU builds (e.g., after `cpu-chatterbox`)
|
||||
@@ -56,24 +70,28 @@ Add `backends/<backend-name>` to the `.NOTPARALLEL` line (around line 2) to prev
|
||||
|
||||
**Step 4b: Add to `prepare-test-extra`**
|
||||
|
||||
Add the backend to the `prepare-test-extra` target (around line 312) to prepare it for testing:
|
||||
Add the backend to the `prepare-test-extra` target to prepare it for testing. Use the path matching your language bucket (`backend/python/`, `backend/go/`, `backend/rust/`, …):
|
||||
|
||||
```makefile
|
||||
prepare-test-extra: protogen-python
|
||||
...
|
||||
$(MAKE) -C backend/python/<backend-name>
|
||||
$(MAKE) -C backend/<lang>/<backend-name>
|
||||
```
|
||||
|
||||
For Rust backends the target is usually the crate build target itself (e.g. `$(MAKE) -C backend/rust/<backend-name> <backend-name>-grpc`) so the binary is in place before `test` runs.
|
||||
|
||||
**Step 4c: Add to `test-extra`**
|
||||
|
||||
Add the backend to the `test-extra` target (around line 319) to run its tests:
|
||||
Add the backend to the `test-extra` target to run its tests — applies to Go and Rust backends too, not only Python:
|
||||
|
||||
```makefile
|
||||
test-extra: prepare-test-extra
|
||||
...
|
||||
$(MAKE) -C backend/python/<backend-name> test
|
||||
$(MAKE) -C backend/<lang>/<backend-name> test
|
||||
```
|
||||
|
||||
Each backend's own `Makefile` should define a `test` target so this line works regardless of language. Integration tests that need large model downloads should be gated behind an env var (see `backend/rust/kokoros/`'s `KOKOROS_MODEL_PATH` pattern) so CI only runs unit tests.
|
||||
|
||||
**Step 4d: Add Backend Definition**
|
||||
|
||||
Add a backend definition variable in the backend definitions section (around line 428-457). The format depends on the backend type:
|
||||
@@ -93,6 +111,13 @@ BACKEND_<BACKEND_NAME> = <backend-name>|python|./backend|false|true
|
||||
BACKEND_<BACKEND_NAME> = <backend-name>|golang|.|false|true
|
||||
```
|
||||
|
||||
**For Rust backends**:
|
||||
```makefile
|
||||
BACKEND_<BACKEND_NAME> = <backend-name>|rust|.|false|true
|
||||
```
|
||||
|
||||
The language field (`python`/`golang`/`rust`/…) must match a `backend/Dockerfile.<lang>` file.
|
||||
|
||||
**Step 4e: Generate Docker Build Target**
|
||||
|
||||
Add an eval call to generate the docker-build target (around line 480-501):
|
||||
@@ -129,6 +154,53 @@ After adding a new backend, verify:
|
||||
- [ ] No Makefile syntax errors (check with linter)
|
||||
- [ ] Follows the same pattern as similar backends (e.g., if it's a transcription backend, follow `faster-whisper` pattern)
|
||||
|
||||
## Bundling runtime shared libraries (`package.sh`)
|
||||
|
||||
The final `Dockerfile.python` stage is `FROM scratch` — there is no system `libc`, no `apt`, no fallback library path. Only files explicitly copied from the builder stage end up in the backend image. That means any runtime `dlopen` your backend (or its Python deps) needs **must** be packaged into `${BACKEND}/lib/`.
|
||||
|
||||
Pattern:
|
||||
|
||||
1. Make sure the library is installed in the builder stage of `backend/Dockerfile.python` (add it to the top-level `apt-get install`).
|
||||
2. Drop a `package.sh` in your backend directory that copies the library — and its soname symlinks — into `$(dirname $0)/lib`. See `backend/python/vllm/package.sh` for a reference implementation that walks `/usr/lib/x86_64-linux-gnu`, `/usr/lib/aarch64-linux-gnu`, etc.
|
||||
3. `Dockerfile.python` already runs `package.sh` automatically if it exists, after `package-gpu-libs.sh`.
|
||||
4. `libbackend.sh` automatically prepends `${EDIR}/lib` to `LD_LIBRARY_PATH` at run time, so anything packaged this way is found by `dlopen`.
|
||||
|
||||
How to find missing libs: when a Python module silently fails to register torch ops or you see `AttributeError: '_OpNamespace' '...' object has no attribute '...'`, run the backend image's Python with `LD_DEBUG=libs` to see which `dlopen` failed. The filename in the error message (e.g. `libnuma.so.1`) is what you need to package.
|
||||
|
||||
To verify packaging works without trusting the host:
|
||||
|
||||
```bash
|
||||
make docker-build-<backend>
|
||||
CID=$(docker create --entrypoint=/run.sh local-ai-backend:<backend>)
|
||||
docker cp $CID:/lib /tmp/check && docker rm $CID
|
||||
ls /tmp/check # expect the bundled .so files + symlinks
|
||||
```
|
||||
|
||||
Then boot it inside a fresh `ubuntu:24.04` (which intentionally does *not* have the lib installed) to confirm it actually loads from the backend dir.
|
||||
|
||||
## Importer integration
|
||||
|
||||
When you add a new backend, you MUST also make it importable via the model import form (`/import-model`). The import form dropdown is sourced dynamically from `GET /backends/known` — it reads the importer registry at `core/gallery/importers/importers.go`, so the steps below are the ONLY way to make your backend show up.
|
||||
|
||||
Required steps:
|
||||
|
||||
1. **If your backend has unambiguous detection signals** (unique file extension, HF `pipeline_tag`, unique repo name pattern, unique artefact like `modules.json`):
|
||||
- Create an importer file at `core/gallery/importers/<backend>.go` following the Match/Import pattern in `llama-cpp.go`.
|
||||
- Register it in `importers.go:defaultImporters` in **specificity order** — more specific detectors must appear BEFORE more generic ones (e.g. `sentencetransformers` before `transformers`, `stablediffusion-ggml` before `llama-cpp`, `vllm-omni` before `vllm`). First match wins.
|
||||
2. **If your backend is a drop-in replacement** (same artefacts as another backend, e.g. `ik-llama-cpp` and `turboquant` both consume GGUF the same way `llama-cpp` does):
|
||||
- Do NOT create a new importer. Extend the existing importer's `Import()` to swap the emitted `backend:` field when `preferences.backend` matches. See `llama-cpp.go` for the pattern.
|
||||
3. **If your backend has no reliable auto-detect signal** (preference-only — e.g. `sglang`, `tinygrad`, `whisperx`):
|
||||
- Do NOT create an importer. Instead add the backend name to the curated pref-only slice in `core/http/endpoints/localai/backend.go` that feeds `/backends/known`. A single line addition.
|
||||
4. **Always** add a table-driven test in `core/gallery/importers/importers_test.go` (Ginkgo/Gomega):
|
||||
- Use a real public HuggingFace repo URI as the test fixture (existing tests already hit the live HF API — follow that pattern).
|
||||
- Cover detection (auto-match without preferences), preference-override (explicit `backend:` in preferences wins), and — if the backend's modality has a common `pipeline_tag` but ambiguous artefacts — an ambiguity test asserting `errors.Is(err, importers.ErrAmbiguousImport)`.
|
||||
|
||||
Rules of thumb:
|
||||
|
||||
- When in doubt, lean pref-only. A wrong auto-detect is worse than a forced preference.
|
||||
- Never silently emit a modality mismatch (e.g. emit `llama-cpp` for a TTS repo because `.gguf` is present). Return `ErrAmbiguousImport` instead.
|
||||
- Registration order is the single most common source of bugs. Check by running `go test ./core/gallery/importers/...` — the existing suite will fail if you've shadowed a pre-existing detector.
|
||||
|
||||
## 6. Example: Adding a Python Backend
|
||||
|
||||
For reference, when `moonshine` was added:
|
||||
|
||||
121
.agents/ai-coding-assistants.md
Normal file
121
.agents/ai-coding-assistants.md
Normal file
@@ -0,0 +1,121 @@
|
||||
# AI Coding Assistants
|
||||
|
||||
This document provides guidance for AI tools and developers using AI
|
||||
assistance when contributing to LocalAI.
|
||||
|
||||
**LocalAI follows the same guidelines as the Linux kernel project for
|
||||
AI-assisted contributions.** See the upstream policy here:
|
||||
<https://docs.kernel.org/process/coding-assistants.html>
|
||||
|
||||
The rules below mirror that policy, adapted to LocalAI's license and
|
||||
project layout. If anything is unclear, the kernel document is the
|
||||
authoritative reference for intent.
|
||||
|
||||
AI tools helping with LocalAI development should follow the standard
|
||||
project development process:
|
||||
|
||||
- [CONTRIBUTING.md](../CONTRIBUTING.md) — development workflow, commit
|
||||
conventions, and PR guidelines
|
||||
- [.agents/coding-style.md](coding-style.md) — code style, editorconfig,
|
||||
logging, and documentation conventions
|
||||
- [.agents/building-and-testing.md](building-and-testing.md) — build and
|
||||
test procedures
|
||||
|
||||
## Licensing and Legal Requirements
|
||||
|
||||
All contributions must comply with LocalAI's licensing requirements:
|
||||
|
||||
- LocalAI is licensed under the **MIT License** — see the [LICENSE](../LICENSE)
|
||||
file
|
||||
- New source files should use the SPDX license identifier `MIT` where
|
||||
applicable to the file type
|
||||
- Contributions must be compatible with the MIT License and must not
|
||||
introduce code under incompatible licenses (e.g., GPL) without an
|
||||
explicit discussion with maintainers
|
||||
|
||||
## Signed-off-by and Developer Certificate of Origin
|
||||
|
||||
Only humans can certify the Developer Certificate of Origin (DCO). AI
|
||||
agents MUST NOT invent or guess a human identity for `Signed-off-by` —
|
||||
doing so forges the DCO certification.
|
||||
|
||||
However, when a human operator explicitly directs the AI to commit on
|
||||
their behalf, the AI is acting as a typing tool — no different from an
|
||||
editor macro or `git commit -s`. In that case the AI SHOULD add
|
||||
`Signed-off-by:` using the **configured `user.name` / `user.email`** of
|
||||
the current git repository (i.e. the operator's own identity). The
|
||||
resulting trailer is the operator's signature; they take responsibility
|
||||
for it by reviewing and pushing the commit. The AI MUST NOT use any
|
||||
other identity and MUST NOT add its own name to the sign-off.
|
||||
|
||||
When running `git commit`, prefer `git commit --signoff` (or `-s`) so
|
||||
the trailer is emitted by git itself from the configured identity,
|
||||
rather than hand-writing it in a heredoc — this guarantees the sign-off
|
||||
matches whatever identity the operator is currently using.
|
||||
|
||||
The human submitter remains responsible for:
|
||||
|
||||
- Reviewing all AI-generated code before it's pushed or merged
|
||||
- Ensuring compliance with licensing requirements
|
||||
- Taking full responsibility for the contribution
|
||||
|
||||
AI agents MUST NOT add `Co-Authored-By` trailers for themselves. A human
|
||||
reviewer owns the contribution; the AI's involvement is recorded via
|
||||
`Assisted-by` (see below).
|
||||
|
||||
## Attribution
|
||||
|
||||
When AI tools contribute to LocalAI development, proper attribution helps
|
||||
track the evolving role of AI in the development process. Contributions
|
||||
should include an `Assisted-by` tag in the commit message trailer in the
|
||||
following format:
|
||||
|
||||
```
|
||||
Assisted-by: AGENT_NAME:MODEL_VERSION [TOOL1] [TOOL2]
|
||||
```
|
||||
|
||||
Where:
|
||||
|
||||
- `AGENT_NAME` — name of the AI tool or framework (e.g., `Claude`,
|
||||
`Copilot`, `Cursor`)
|
||||
- `MODEL_VERSION` — specific model version used (e.g.,
|
||||
`claude-opus-4-7`, `gpt-5`)
|
||||
- `[TOOL1] [TOOL2]` — optional specialized analysis tools invoked by the
|
||||
agent (e.g., `golangci-lint`, `staticcheck`, `go vet`)
|
||||
|
||||
Basic development tools (git, go, make, editors) should **not** be listed.
|
||||
|
||||
### Example
|
||||
|
||||
```
|
||||
fix(llama-cpp): handle empty tool call arguments
|
||||
|
||||
Previously the parser panicked when the model returned a tool call with
|
||||
an empty arguments object. Fall back to an empty JSON object in that
|
||||
case so downstream consumers receive a valid payload.
|
||||
|
||||
Assisted-by: Claude:claude-opus-4-7 golangci-lint
|
||||
Signed-off-by: Jane Developer <jane@example.com>
|
||||
```
|
||||
|
||||
The `Signed-off-by` line uses Jane's own identity because Jane is the
|
||||
submitter operating the AI. If Jane asks Claude to create the commit via
|
||||
`git commit -s`, git emits that exact trailer from Jane's configured
|
||||
identity — no separate human step is needed beyond Jane reviewing the
|
||||
diff before pushing.
|
||||
|
||||
## Scope and Responsibility
|
||||
|
||||
Using an AI assistant does not reduce the contributor's responsibility.
|
||||
The human submitter must:
|
||||
|
||||
- Understand every line that lands in the PR
|
||||
- Verify that generated code compiles, passes tests, and follows the
|
||||
project style
|
||||
- Confirm that any referenced APIs, flags, or file paths actually exist
|
||||
in the current tree (AI models may hallucinate identifiers)
|
||||
- Not submit AI output verbatim without review
|
||||
|
||||
Reviewers may ask for clarification on any change regardless of how it
|
||||
was produced. "An AI wrote it" is not an acceptable answer to a design
|
||||
question.
|
||||
@@ -2,6 +2,8 @@
|
||||
|
||||
This guide covers how to add new API endpoints and properly integrate them with the auth/permissions system.
|
||||
|
||||
> **Before you ship a new endpoint or capability surface**, re-read the [checklist at the bottom of this file](#checklist). LocalAI advertises its feature surface in several independent places — miss any one of them and clients/admins/UI won't know the endpoint exists.
|
||||
|
||||
## Architecture overview
|
||||
|
||||
Authentication and authorization flow through three layers:
|
||||
@@ -234,6 +236,66 @@ Use these HTTP status codes:
|
||||
|
||||
If your endpoint should be tracked for usage (token counts, request counts), add the `usageMiddleware` to its middleware chain. See `core/http/middleware/usage.go` and how it's applied in `routes/openai.go`.
|
||||
|
||||
## Advertising surfaces — where to register a new capability
|
||||
|
||||
Beyond routing and auth, LocalAI publishes its capability surface in **four independent places**. When you add an endpoint — especially one introducing a net-new capability like a new media type or a new auth-gated feature — you must update every relevant surface. These aren't optional: missing them means the endpoint works but is invisible to clients, admins, and the UI.
|
||||
|
||||
### 1. Swagger `@Tags` annotation (mandatory)
|
||||
|
||||
Every handler needs a swagger block so the endpoint appears in `/swagger/index.html` and in the `/api/instructions` output. The `@Tags` value is what groups the endpoint into a capability area:
|
||||
|
||||
```go
|
||||
// MyEndpoint does X.
|
||||
// @Summary Do X.
|
||||
// @Tags my-capability
|
||||
// @Param request body schema.MyRequest true "payload"
|
||||
// @Success 200 {object} schema.MyResponse "Response"
|
||||
// @Router /v1/my-endpoint [post]
|
||||
func MyEndpoint(...) echo.HandlerFunc { ... }
|
||||
```
|
||||
|
||||
Use an existing tag when the endpoint extends an existing area (e.g. `audio`, `images`, `face-recognition`). Create a new tag only when the endpoint introduces a genuinely new capability surface — and in that case, also register it in step 2.
|
||||
|
||||
After adding endpoints, regenerate the embedded spec so the runtime serves it:
|
||||
|
||||
```bash
|
||||
make protogen-go # ensures gRPC codegen is fresh first
|
||||
make swagger # regenerates swagger/swagger.json
|
||||
```
|
||||
|
||||
### 2. `/api/instructions` registry (for new capability areas)
|
||||
|
||||
`core/http/endpoints/localai/api_instructions.go` defines `instructionDefs` — a lightweight, machine-readable index of capability areas that groups swagger endpoints by tag. It's the primary discovery surface for agents and SDKs ("what can this server do?").
|
||||
|
||||
**When to update:** only when adding a new capability area (a new swagger tag). Existing-tag additions automatically surface without any change here.
|
||||
|
||||
Add an entry to `instructionDefs`:
|
||||
|
||||
```go
|
||||
{
|
||||
Name: "my-capability", // URL segment at /api/instructions/my-capability
|
||||
Description: "Short sentence describing the capability",
|
||||
Tags: []string{"my-capability"}, // must match swagger @Tags
|
||||
Intro: "Optional gotcha/context that isn't in the swagger descriptions (caveats, defaults, cross-references to other endpoints).",
|
||||
},
|
||||
```
|
||||
|
||||
Also bump the expected-length count in `api_instructions_test.go` and add the name to the `ContainElements` assertion.
|
||||
|
||||
### 3. `capabilities.js` symbol (for new model-config FLAG_* flags)
|
||||
|
||||
If your feature needs a new `FLAG_*` usecase flag in `core/config/model_config.go` (so users can filter gallery models by it, and so `/v1/models` surfaces it), also declare the matching symbol in `core/http/react-ui/src/utils/capabilities.js`:
|
||||
|
||||
```js
|
||||
export const CAP_MY_CAPABILITY = 'FLAG_MY_CAPABILITY'
|
||||
```
|
||||
|
||||
React pages that want to filter the ModelSelector by capability import this symbol. Declare it even if you're not building the UI page yet — the declaration keeps the Go/JS vocabularies in sync.
|
||||
|
||||
### 4. `docs/content/` (user-facing documentation)
|
||||
|
||||
A new capability deserves its own page under `docs/content/features/`, plus cross-links from related features and an entry in `docs/content/whats-new.md`. See the pattern used by `face-recognition.md` / `object-detection.md`.
|
||||
|
||||
## Path protection rules
|
||||
|
||||
The global auth middleware classifies paths as API paths or non-API paths:
|
||||
@@ -248,12 +310,23 @@ If you add endpoints under a new top-level path prefix, add it to `isAPIPath()`
|
||||
|
||||
When adding a new endpoint:
|
||||
|
||||
**Routing & auth**
|
||||
- [ ] Handler in `core/http/endpoints/`
|
||||
- [ ] Route registered in appropriate `core/http/routes/` file
|
||||
- [ ] Auth level chosen: public / standard / admin / feature-gated
|
||||
- [ ] If feature-gated: constant in `permissions.go`, metadata in `features.go`, middleware in `app.go`
|
||||
- [ ] Entry added to `RouteFeatureRegistry` in `core/http/auth/features.go` (one row per route/method — all /v1/* routes gate through this, not per-route middleware)
|
||||
- [ ] If new feature: constant in `permissions.go`, added to the right slice (`APIFeatures` default-ON / `AgentFeatures` default-OFF), metadata in `features.go` `*FeatureMetas()`
|
||||
- [ ] If feature uses group middleware: wired in `core/http/app.go` and passed to the route registration function
|
||||
- [ ] If new path prefix: added to `isAPIPath()` in `middleware.go`
|
||||
- [ ] If OpenAI-compatible: entry in `RouteFeatureRegistry`
|
||||
- [ ] If token-counting: `usageMiddleware` added to middleware chain
|
||||
- [ ] Error responses use `schema.ErrorResponse` format
|
||||
|
||||
**Advertising surfaces (easy to miss — see the [Advertising surfaces](#advertising-surfaces--where-to-register-a-new-capability) section)**
|
||||
- [ ] Swagger block on the handler: `@Summary`, `@Tags`, `@Param`, `@Success`, `@Router`
|
||||
- [ ] If new capability area (new swagger tag): entry in `instructionDefs` in `core/http/endpoints/localai/api_instructions.go` + test count bumped in `api_instructions_test.go`
|
||||
- [ ] If new `FLAG_*` usecase flag: matching `CAP_*` symbol exported from `core/http/react-ui/src/utils/capabilities.js`
|
||||
- [ ] `docs/content/features/<feature>.md` created; cross-links from related feature pages; entry in `docs/content/whats-new.md`
|
||||
|
||||
**Quality**
|
||||
- [ ] Error responses use `schema.ErrorResponse` format (or `echo.NewHTTPError` with a mapped gRPC status — see the `mapBackendError` helper in `core/http/endpoints/localai/images.go`)
|
||||
- [ ] Tests cover both authenticated and unauthenticated access
|
||||
- [ ] Swagger regenerated (`make swagger`) if you changed any `@Router`/`@Tags`/`@Param` annotation
|
||||
|
||||
@@ -42,6 +42,12 @@ trim_trailing_whitespace = false
|
||||
|
||||
Use `github.com/mudler/xlog` for logging which has the same API as slog.
|
||||
|
||||
## Go tests
|
||||
|
||||
All Go tests — including backend tests — must use [Ginkgo](https://onsi.github.io/ginkgo/) (v2) with Gomega matchers, not the stdlib `testing` package with `t.Run` / `t.Errorf`. A test file should register a suite with `RegisterFailHandler(Fail)` in a `TestXxx(t *testing.T)` bootstrap and use `Describe`/`Context`/`It` blocks for the actual cases. Look at any existing `*_test.go` under `core/` or `pkg/` for a template.
|
||||
|
||||
Do not mix styles within a package. If you are extending tests in a package that already uses Ginkgo, keep using Ginkgo. If you find stdlib-style Go tests in the tree, treat them as tech debt to be migrated rather than as a pattern to follow.
|
||||
|
||||
## Documentation
|
||||
|
||||
The project documentation is located in `docs/content`. When adding new features or changing existing functionality, it is crucial to update the documentation to reflect these changes. This helps users understand how to use the new capabilities and ensures the documentation stays relevant.
|
||||
|
||||
115
.agents/vllm-backend.md
Normal file
115
.agents/vllm-backend.md
Normal file
@@ -0,0 +1,115 @@
|
||||
# Working on the vLLM Backend
|
||||
|
||||
The vLLM backend lives at `backend/python/vllm/backend.py` (async gRPC) and the multimodal variant at `backend/python/vllm-omni/backend.py` (sync gRPC). Both wrap vLLM's `AsyncLLMEngine` / `Omni` and translate the LocalAI gRPC `PredictOptions` into vLLM `SamplingParams` + outputs into `Reply.chat_deltas`.
|
||||
|
||||
This file captures the non-obvious bits — most of the bring-up was a single PR (`feat/vllm-parity`) and the things below are easy to get wrong.
|
||||
|
||||
## Tool calling and reasoning use vLLM's *native* parsers
|
||||
|
||||
Do not write regex-based tool-call extractors for vLLM. vLLM ships:
|
||||
|
||||
- `vllm.tool_parsers.ToolParserManager` — 50+ registered parsers (`hermes`, `llama3_json`, `llama4_pythonic`, `mistral`, `qwen3_xml`, `deepseek_v3`, `granite4`, `openai`, `kimi_k2`, `glm45`, …)
|
||||
- `vllm.reasoning.ReasoningParserManager` — 25+ registered parsers (`deepseek_r1`, `qwen3`, `mistral`, `gemma4`, …)
|
||||
|
||||
Both can be used standalone: instantiate with a tokenizer, call `extract_tool_calls(text, request=None)` / `extract_reasoning(text, request=None)`. The backend stores the parser *classes* on `self.tool_parser_cls` / `self.reasoning_parser_cls` at LoadModel time and instantiates them per request.
|
||||
|
||||
**Selection:** vLLM does *not* auto-detect parsers from model name — neither does the LocalAI backend. The user (or `core/config/hooks_vllm.go`) must pick one and pass it via `Options[]`:
|
||||
|
||||
```yaml
|
||||
options:
|
||||
- tool_parser:hermes
|
||||
- reasoning_parser:qwen3
|
||||
```
|
||||
|
||||
Auto-defaults for known model families live in `core/config/parser_defaults.json` and are applied:
|
||||
- at gallery import time by `core/gallery/importers/vllm.go`
|
||||
- at model load time by the `vllm` / `vllm-omni` backend hook in `core/config/hooks_vllm.go`
|
||||
|
||||
User-supplied `tool_parser:`/`reasoning_parser:` in the config wins over defaults — the hook checks for existing entries before appending.
|
||||
|
||||
**When to update `parser_defaults.json`:** any time vLLM ships a new tool or reasoning parser, or you onboard a new model family that LocalAI users will pull from HuggingFace. The file is keyed by *family pattern* matched against `normalizeModelID(cfg.Model)` (lowercase, org-prefix stripped, `_`→`-`). Patterns are checked **longest-first** — keep `qwen3.5` before `qwen3`, `llama-3.3` before `llama-3`, etc., or the wrong family wins. Add a covering test in `core/config/hooks_test.go`.
|
||||
|
||||
**Sister file — `core/config/inference_defaults.json`:** same pattern but for sampling parameters (temperature, top_p, top_k, min_p, repeat_penalty, presence_penalty). Loaded by `core/config/inference_defaults.go` and applied by `ApplyInferenceDefaults()`. The schema is `map[string]float64` only — *strings don't fit*, which is why parser defaults needed their own JSON file. The inference file is **auto-generated from unsloth** via `go generate ./core/config/` (see `core/config/gen_inference_defaults/`) — don't hand-edit it; instead update the upstream source or regenerate. Both files share `normalizeModelID()` and the longest-first pattern ordering.
|
||||
|
||||
**Constructor compatibility gotcha:** the abstract `ToolParser.__init__` accepts `tools=`, but several concrete parsers (Hermes2ProToolParser, etc.) override `__init__` and *only* accept `tokenizer`. Always:
|
||||
|
||||
```python
|
||||
try:
|
||||
tp = self.tool_parser_cls(self.tokenizer, tools=tools)
|
||||
except TypeError:
|
||||
tp = self.tool_parser_cls(self.tokenizer)
|
||||
```
|
||||
|
||||
## ChatDelta is the streaming contract
|
||||
|
||||
The Go side (`core/backend/llm.go`, `pkg/functions/chat_deltas.go`) consumes `Reply.chat_deltas` to assemble the OpenAI response. For tool calls to surface in `chat/completions`, the Python backend **must** populate `Reply.chat_deltas[].tool_calls` with `ToolCallDelta{index, id, name, arguments}`. Returning the raw `<tool_call>...</tool_call>` text in `Reply.message` is *not* enough — the Go regex fallback exists for llama.cpp, not for vllm.
|
||||
|
||||
Same story for `reasoning_content` — emit it on `ChatDelta.reasoning_content`, not as part of `content`.
|
||||
|
||||
## Message conversion to chat templates
|
||||
|
||||
`tokenizer.apply_chat_template()` expects a list of dicts, not proto Messages. The shared helper in `backend/python/common/vllm_utils.py` (`messages_to_dicts`) handles the mapping including:
|
||||
|
||||
- `tool_call_id` and `name` for `role="tool"` messages
|
||||
- `tool_calls` JSON-string field → parsed Python list for `role="assistant"`
|
||||
- `reasoning_content` for thinking models
|
||||
|
||||
Pass `tools=json.loads(request.Tools)` and (when `request.Metadata.get("enable_thinking") == "true"`) `enable_thinking=True` to `apply_chat_template`. Wrap in `try/except TypeError` because not every tokenizer template accepts those kwargs.
|
||||
|
||||
## CPU support and the SIMD/library minefield
|
||||
|
||||
vLLM publishes prebuilt CPU wheels at `https://github.com/vllm-project/vllm/releases/...`. The pin lives in `backend/python/vllm/requirements-cpu-after.txt`.
|
||||
|
||||
**Version compatibility — important:** newer vllm CPU wheels (≥ 0.15) declare `torch==2.10.0+cpu` as a hard dep, but `torch==2.10.0` only exists on the PyTorch test channel and pulls in an incompatible `torchvision`. Stay on **`vllm 0.14.1+cpu` + `torch 2.9.1+cpu`** until both upstream catch up. Bumping requires verifying torchvision/torchaudio match.
|
||||
|
||||
`requirements-cpu.txt` uses `--extra-index-url https://download.pytorch.org/whl/cpu`. `install.sh` adds `--index-strategy=unsafe-best-match` for the `cpu` profile so uv resolves transformers/vllm from PyPI while pulling torch from the PyTorch index.
|
||||
|
||||
**SIMD baseline:** the prebuilt CPU wheel is compiled with AVX-512 VNNI/BF16. On a CPU without those instructions, importing `vllm.model_executor.models.registry` SIGILLs at `_run_in_subprocess` time during model inspection. There is no runtime flag to disable it. Workarounds:
|
||||
|
||||
1. **Run on a host with the right SIMD baseline** (default — fast)
|
||||
2. **Build from source** with `FROM_SOURCE=true` env var. Plumbing exists end-to-end:
|
||||
- `install.sh` hides `requirements-cpu-after.txt`, runs `installRequirements` for the base deps, then clones vllm and `VLLM_TARGET_DEVICE=cpu uv pip install --no-deps .`
|
||||
- `backend/Dockerfile.python` declares `ARG FROM_SOURCE` + `ENV FROM_SOURCE`
|
||||
- `Makefile` `docker-build-backend` macro forwards `--build-arg FROM_SOURCE=$(FROM_SOURCE)` when set
|
||||
- Source build takes 30–50 minutes — too slow for per-PR CI but fine for local.
|
||||
|
||||
**Runtime shared libraries:** vLLM's `vllm._C` extension `dlopen`s `libnuma.so.1` at import time. If missing, the C extension silently fails and `torch.ops._C_utils.init_cpu_threads_env` is never registered → `EngineCore` crashes on `init_device` with:
|
||||
|
||||
```
|
||||
AttributeError: '_OpNamespace' '_C_utils' object has no attribute 'init_cpu_threads_env'
|
||||
```
|
||||
|
||||
`backend/python/vllm/package.sh` bundles `libnuma.so.1` and `libgomp.so.1` into `${BACKEND}/lib/`, which `libbackend.sh` adds to `LD_LIBRARY_PATH` at run time. The builder stage in `backend/Dockerfile.python` installs `libnuma1`/`libgomp1` so package.sh has something to copy. Do *not* assume the production host has these — backend images are `FROM scratch`.
|
||||
|
||||
## Backend hook system (`core/config/backend_hooks.go`)
|
||||
|
||||
Per-backend defaults that used to be hardcoded in `ModelConfig.Prepare()` now live in `core/config/hooks_*.go` files and self-register via `init()`:
|
||||
|
||||
- `hooks_llamacpp.go` → GGUF metadata parsing, context size, GPU layers, jinja template
|
||||
- `hooks_vllm.go` → tool/reasoning parser auto-selection from `parser_defaults.json`
|
||||
|
||||
Hook keys:
|
||||
- `"llama-cpp"`, `"vllm"`, `"vllm-omni"`, … — backend-specific
|
||||
- `""` — runs only when `cfg.Backend` is empty (auto-detect case)
|
||||
- `"*"` — global catch-all, runs for every backend before specific hooks
|
||||
|
||||
Multiple hooks per key are supported and run in registration order. Adding a new backend default:
|
||||
|
||||
```go
|
||||
// core/config/hooks_<backend>.go
|
||||
func init() {
|
||||
RegisterBackendHook("<backend>", myDefaults)
|
||||
}
|
||||
func myDefaults(cfg *ModelConfig, modelPath string) {
|
||||
// only fill in fields the user didn't set
|
||||
}
|
||||
```
|
||||
|
||||
## The `Messages.ToProto()` fields you need to set
|
||||
|
||||
`core/schema/message.go:ToProto()` must serialize:
|
||||
- `ToolCallID` → `proto.Message.ToolCallId` (for `role="tool"` messages — links result back to the call)
|
||||
- `Reasoning` → `proto.Message.ReasoningContent`
|
||||
- `ToolCalls` → `proto.Message.ToolCalls` (JSON-encoded string)
|
||||
|
||||
These were originally not serialized and tool-calling conversations broke silently — the C++ llama.cpp backend reads them but always got empty strings. Any new field added to `schema.Message` *and* `proto.Message` needs a matching line in `ToProto()`.
|
||||
446
.github/gallery-agent/agent.go
vendored
446
.github/gallery-agent/agent.go
vendored
@@ -1,446 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"regexp"
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"github.com/ghodss/yaml"
|
||||
hfapi "github.com/mudler/LocalAI/pkg/huggingface-api"
|
||||
"github.com/mudler/cogito"
|
||||
"github.com/mudler/cogito/clients"
|
||||
"github.com/mudler/cogito/structures"
|
||||
"github.com/sashabaranov/go-openai/jsonschema"
|
||||
)
|
||||
|
||||
var (
|
||||
openAIModel = os.Getenv("OPENAI_MODEL")
|
||||
openAIKey = os.Getenv("OPENAI_KEY")
|
||||
openAIBaseURL = os.Getenv("OPENAI_BASE_URL")
|
||||
galleryIndexPath = os.Getenv("GALLERY_INDEX_PATH")
|
||||
//defaultclient
|
||||
llm = clients.NewOpenAILLM(openAIModel, openAIKey, openAIBaseURL)
|
||||
)
|
||||
|
||||
// cleanTextContent removes trailing spaces, tabs, and normalizes line endings
|
||||
// to prevent YAML linting issues like trailing spaces and multiple empty lines
|
||||
func cleanTextContent(text string) string {
|
||||
lines := strings.Split(text, "\n")
|
||||
var cleanedLines []string
|
||||
var prevEmpty bool
|
||||
for _, line := range lines {
|
||||
// Remove all trailing whitespace (spaces, tabs, etc.)
|
||||
trimmed := strings.TrimRight(line, " \t\r")
|
||||
// Avoid multiple consecutive empty lines
|
||||
if trimmed == "" {
|
||||
if !prevEmpty {
|
||||
cleanedLines = append(cleanedLines, "")
|
||||
}
|
||||
prevEmpty = true
|
||||
} else {
|
||||
cleanedLines = append(cleanedLines, trimmed)
|
||||
prevEmpty = false
|
||||
}
|
||||
}
|
||||
// Remove trailing empty lines from the result
|
||||
result := strings.Join(cleanedLines, "\n")
|
||||
return stripThinkingTags(strings.TrimRight(result, "\n"))
|
||||
}
|
||||
|
||||
type galleryModel struct {
|
||||
Name string `yaml:"name"`
|
||||
Urls []string `yaml:"urls"`
|
||||
}
|
||||
|
||||
// isModelExisting checks if a specific model ID exists in the gallery using text search
|
||||
func isModelExisting(modelID string) (bool, error) {
|
||||
indexPath := getGalleryIndexPath()
|
||||
content, err := os.ReadFile(indexPath)
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("failed to read %s: %w", indexPath, err)
|
||||
}
|
||||
|
||||
var galleryModels []galleryModel
|
||||
|
||||
err = yaml.Unmarshal(content, &galleryModels)
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("failed to unmarshal %s: %w", indexPath, err)
|
||||
}
|
||||
|
||||
for _, galleryModel := range galleryModels {
|
||||
if slices.Contains(galleryModel.Urls, modelID) {
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// filterExistingModels removes models that already exist in the gallery
|
||||
func filterExistingModels(models []ProcessedModel) ([]ProcessedModel, error) {
|
||||
var filteredModels []ProcessedModel
|
||||
for _, model := range models {
|
||||
exists, err := isModelExisting(model.ModelID)
|
||||
if err != nil {
|
||||
fmt.Printf("Error checking if model %s exists: %v, skipping\n", model.ModelID, err)
|
||||
continue
|
||||
}
|
||||
|
||||
if !exists {
|
||||
filteredModels = append(filteredModels, model)
|
||||
} else {
|
||||
fmt.Printf("Skipping existing model: %s\n", model.ModelID)
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Printf("Filtered out %d existing models, %d new models remaining\n",
|
||||
len(models)-len(filteredModels), len(filteredModels))
|
||||
|
||||
return filteredModels, nil
|
||||
}
|
||||
|
||||
// getGalleryIndexPath returns the gallery index file path, with a default fallback
|
||||
func getGalleryIndexPath() string {
|
||||
if galleryIndexPath != "" {
|
||||
return galleryIndexPath
|
||||
}
|
||||
return "gallery/index.yaml"
|
||||
}
|
||||
|
||||
func stripThinkingTags(content string) string {
|
||||
// Remove content between <thinking> and </thinking> (including multi-line)
|
||||
content = regexp.MustCompile(`(?s)<thinking>.*?</thinking>`).ReplaceAllString(content, "")
|
||||
// Remove content between <think> and </think> (including multi-line)
|
||||
content = regexp.MustCompile(`(?s)<think>.*?</think>`).ReplaceAllString(content, "")
|
||||
// Clean up any extra whitespace
|
||||
content = strings.TrimSpace(content)
|
||||
return content
|
||||
}
|
||||
|
||||
func getRealReadme(ctx context.Context, repository string) (string, error) {
|
||||
// Create a conversation fragment
|
||||
fragment := cogito.NewEmptyFragment().
|
||||
AddMessage("user",
|
||||
`Your task is to get a clear description of a large language model from huggingface by using the provided tool. I will share with you a repository that might be quantized, and as such probably not by the original model author. We need to get the real description of the model, and not the one that might be quantized. You will have to call the tool to get the readme more than once by figuring out from the quantized readme which is the base model readme. This is the repository: `+repository)
|
||||
|
||||
// Execute with tools
|
||||
result, err := cogito.ExecuteTools(llm, fragment,
|
||||
cogito.WithIterations(3),
|
||||
cogito.WithMaxAttempts(3),
|
||||
cogito.DisableSinkState,
|
||||
cogito.WithTools(&HFReadmeTool{client: hfapi.NewClient()}))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
result = result.AddMessage("user", "Describe the model in a clear and concise way that can be shared in a model gallery.")
|
||||
|
||||
// Get a response
|
||||
_, err = llm.Ask(ctx, result)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
content := result.LastMessage().Content
|
||||
return cleanTextContent(content), nil
|
||||
}
|
||||
|
||||
func selectMostInterestingModels(ctx context.Context, searchResult *SearchResult) ([]ProcessedModel, error) {
|
||||
|
||||
if len(searchResult.Models) == 1 {
|
||||
return searchResult.Models, nil
|
||||
}
|
||||
|
||||
// Create a conversation fragment
|
||||
fragment := cogito.NewEmptyFragment().
|
||||
AddMessage("user",
|
||||
`Your task is to analyze a list of AI models and select the most interesting ones for a model gallery. You will be given detailed information about multiple models including their metadata, file information, and README content.
|
||||
|
||||
Consider the following criteria when selecting models:
|
||||
1. Model popularity (download count)
|
||||
2. Model recency (last modified date)
|
||||
3. Model completeness (has preferred model file, README, etc.)
|
||||
4. Model uniqueness (not duplicates or very similar models)
|
||||
5. Model quality (based on README content and description)
|
||||
6. Model utility (practical applications)
|
||||
|
||||
You should select models that would be most valuable for users browsing a model gallery. Prioritize models that are:
|
||||
- Well-documented with clear READMEs
|
||||
- Recently updated
|
||||
- Popular (high download count)
|
||||
- Have the preferred quantization format available
|
||||
- Offer unique capabilities or are from reputable authors
|
||||
|
||||
Return your analysis and selection reasoning.`)
|
||||
|
||||
// Add the search results as context
|
||||
modelsInfo := fmt.Sprintf("Found %d models matching '%s' with quantization preference '%s':\n\n",
|
||||
searchResult.TotalModelsFound, searchResult.SearchTerm, searchResult.Quantization)
|
||||
|
||||
for i, model := range searchResult.Models {
|
||||
modelsInfo += fmt.Sprintf("Model %d:\n", i+1)
|
||||
modelsInfo += fmt.Sprintf(" ID: %s\n", model.ModelID)
|
||||
modelsInfo += fmt.Sprintf(" Author: %s\n", model.Author)
|
||||
modelsInfo += fmt.Sprintf(" Downloads: %d\n", model.Downloads)
|
||||
modelsInfo += fmt.Sprintf(" Last Modified: %s\n", model.LastModified)
|
||||
modelsInfo += fmt.Sprintf(" Files: %d files\n", len(model.Files))
|
||||
|
||||
if model.PreferredModelFile != nil {
|
||||
modelsInfo += fmt.Sprintf(" Preferred Model File: %s (%d bytes)\n",
|
||||
model.PreferredModelFile.Path, model.PreferredModelFile.Size)
|
||||
} else {
|
||||
modelsInfo += " No preferred model file found\n"
|
||||
}
|
||||
|
||||
if model.ReadmeContent != "" {
|
||||
modelsInfo += fmt.Sprintf(" README: %s\n", model.ReadmeContent)
|
||||
}
|
||||
|
||||
if model.ProcessingError != "" {
|
||||
modelsInfo += fmt.Sprintf(" Processing Error: %s\n", model.ProcessingError)
|
||||
}
|
||||
|
||||
modelsInfo += "\n"
|
||||
}
|
||||
|
||||
fragment = fragment.AddMessage("user", modelsInfo)
|
||||
|
||||
fragment = fragment.AddMessage("user", "Based on your analysis, select the top 5 most interesting models and provide a brief explanation for each selection. Also, create a filtered SearchResult with only the selected models. Return just a list of repositories IDs, you will later be asked to output it as a JSON array with the json tool.")
|
||||
|
||||
// Get a response
|
||||
newFragment, err := llm.Ask(ctx, fragment)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
fmt.Println(newFragment.LastMessage().Content)
|
||||
repositories := struct {
|
||||
Repositories []string `json:"repositories"`
|
||||
}{}
|
||||
|
||||
s := structures.Structure{
|
||||
Schema: jsonschema.Definition{
|
||||
Type: jsonschema.Object,
|
||||
AdditionalProperties: false,
|
||||
Properties: map[string]jsonschema.Definition{
|
||||
"repositories": {
|
||||
Type: jsonschema.Array,
|
||||
Items: &jsonschema.Definition{Type: jsonschema.String},
|
||||
Description: "The trending repositories IDs",
|
||||
},
|
||||
},
|
||||
Required: []string{"repositories"},
|
||||
},
|
||||
Object: &repositories,
|
||||
}
|
||||
|
||||
err = newFragment.ExtractStructure(ctx, llm, s)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
filteredModels := []ProcessedModel{}
|
||||
for _, m := range searchResult.Models {
|
||||
if slices.Contains(repositories.Repositories, m.ModelID) {
|
||||
filteredModels = append(filteredModels, m)
|
||||
}
|
||||
}
|
||||
|
||||
return filteredModels, nil
|
||||
}
|
||||
|
||||
// ModelMetadata represents extracted metadata from a model
|
||||
type ModelMetadata struct {
|
||||
Tags []string `json:"tags"`
|
||||
License string `json:"license"`
|
||||
}
|
||||
|
||||
// extractModelMetadata extracts tags and license from model README and documentation
|
||||
func extractModelMetadata(ctx context.Context, model ProcessedModel) ([]string, string, error) {
|
||||
// Create a conversation fragment
|
||||
fragment := cogito.NewEmptyFragment().
|
||||
AddMessage("user",
|
||||
`Your task is to extract metadata from an AI model's README and documentation. You will be provided with:
|
||||
1. Model information (ID, author, description)
|
||||
2. README content
|
||||
|
||||
You need to extract:
|
||||
1. **Tags**: An array of relevant tags that describe the model. Use common tags from the gallery such as:
|
||||
- llm, gguf, gpu, cpu, multimodal, image-to-text, text-to-text, text-to-speech, tts
|
||||
- thinking, reasoning, chat, instruction-tuned, code, vision
|
||||
- Model family names (e.g., llama, qwen, mistral, gemma) if applicable
|
||||
- Any other relevant descriptive tags
|
||||
Select 3-8 most relevant tags.
|
||||
|
||||
2. **License**: The license identifier (e.g., "apache-2.0", "mit", "llama2", "gpl-3.0", "bsd", "cc-by-4.0").
|
||||
If no license is found, return an empty string.
|
||||
|
||||
Return the extracted metadata in a structured format.`)
|
||||
|
||||
// Add model information
|
||||
modelInfo := "Model Information:\n"
|
||||
modelInfo += fmt.Sprintf(" ID: %s\n", model.ModelID)
|
||||
modelInfo += fmt.Sprintf(" Author: %s\n", model.Author)
|
||||
modelInfo += fmt.Sprintf(" Downloads: %d\n", model.Downloads)
|
||||
if model.ReadmeContent != "" {
|
||||
modelInfo += fmt.Sprintf(" README Content:\n%s\n", model.ReadmeContent)
|
||||
} else if model.ReadmeContentPreview != "" {
|
||||
modelInfo += fmt.Sprintf(" README Preview: %s\n", model.ReadmeContentPreview)
|
||||
}
|
||||
|
||||
fragment = fragment.AddMessage("user", modelInfo)
|
||||
fragment = fragment.AddMessage("user", "Extract the tags and license from the model information. Return the metadata as a JSON object with 'tags' (array of strings) and 'license' (string).")
|
||||
|
||||
// Get a response
|
||||
newFragment, err := llm.Ask(ctx, fragment)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
// Extract structured metadata
|
||||
metadata := ModelMetadata{}
|
||||
|
||||
s := structures.Structure{
|
||||
Schema: jsonschema.Definition{
|
||||
Type: jsonschema.Object,
|
||||
AdditionalProperties: false,
|
||||
Properties: map[string]jsonschema.Definition{
|
||||
"tags": {
|
||||
Type: jsonschema.Array,
|
||||
Items: &jsonschema.Definition{Type: jsonschema.String},
|
||||
Description: "Array of relevant tags describing the model",
|
||||
},
|
||||
"license": {
|
||||
Type: jsonschema.String,
|
||||
Description: "License identifier (e.g., apache-2.0, mit, llama2). Empty string if not found.",
|
||||
},
|
||||
},
|
||||
Required: []string{"tags", "license"},
|
||||
},
|
||||
Object: &metadata,
|
||||
}
|
||||
|
||||
err = newFragment.ExtractStructure(ctx, llm, s)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
return metadata.Tags, metadata.License, nil
|
||||
}
|
||||
|
||||
// extractIconFromReadme scans the README content for image URLs and returns the first suitable icon URL found
|
||||
func extractIconFromReadme(readmeContent string) string {
|
||||
if readmeContent == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
// Regular expressions to match image URLs in various formats (case-insensitive)
|
||||
// Match markdown image syntax:  - case insensitive extensions
|
||||
markdownImageRegex := regexp.MustCompile(`(?i)!\[[^\]]*\]\(([^)]+\.(png|jpg|jpeg|svg|webp|gif))\)`)
|
||||
// Match HTML img tags: <img src="url">
|
||||
htmlImageRegex := regexp.MustCompile(`(?i)<img[^>]+src=["']([^"']+\.(png|jpg|jpeg|svg|webp|gif))["']`)
|
||||
// Match plain URLs ending with image extensions
|
||||
plainImageRegex := regexp.MustCompile(`(?i)https?://[^\s<>"']+\.(png|jpg|jpeg|svg|webp|gif)`)
|
||||
|
||||
// Try markdown format first
|
||||
matches := markdownImageRegex.FindStringSubmatch(readmeContent)
|
||||
if len(matches) > 1 && matches[1] != "" {
|
||||
url := strings.TrimSpace(matches[1])
|
||||
// Prefer HuggingFace CDN URLs or absolute URLs
|
||||
if strings.HasPrefix(strings.ToLower(url), "http") {
|
||||
return url
|
||||
}
|
||||
}
|
||||
|
||||
// Try HTML img tags
|
||||
matches = htmlImageRegex.FindStringSubmatch(readmeContent)
|
||||
if len(matches) > 1 && matches[1] != "" {
|
||||
url := strings.TrimSpace(matches[1])
|
||||
if strings.HasPrefix(strings.ToLower(url), "http") {
|
||||
return url
|
||||
}
|
||||
}
|
||||
|
||||
// Try plain URLs
|
||||
matches = plainImageRegex.FindStringSubmatch(readmeContent)
|
||||
if len(matches) > 0 {
|
||||
url := strings.TrimSpace(matches[0])
|
||||
if strings.HasPrefix(strings.ToLower(url), "http") {
|
||||
return url
|
||||
}
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
// getHuggingFaceAvatarURL attempts to get the HuggingFace avatar URL for a user
|
||||
func getHuggingFaceAvatarURL(author string) string {
|
||||
if author == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
// Try to fetch user info from HuggingFace API
|
||||
// HuggingFace API endpoint: https://huggingface.co/api/users/{username}
|
||||
baseURL := "https://huggingface.co"
|
||||
userURL := fmt.Sprintf("%s/api/users/%s", baseURL, author)
|
||||
|
||||
req, err := http.NewRequest("GET", userURL, nil)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
client := &http.Client{}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return ""
|
||||
}
|
||||
|
||||
// Parse the response to get avatar URL
|
||||
var userInfo map[string]any
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(body, &userInfo); err != nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
// Try to extract avatar URL from response
|
||||
if avatar, ok := userInfo["avatarUrl"].(string); ok && avatar != "" {
|
||||
return avatar
|
||||
}
|
||||
if avatar, ok := userInfo["avatar"].(string); ok && avatar != "" {
|
||||
return avatar
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
// extractModelIcon extracts icon URL from README or falls back to HuggingFace avatar
|
||||
func extractModelIcon(model ProcessedModel) string {
|
||||
// First, try to extract icon from README
|
||||
if icon := extractIconFromReadme(model.ReadmeContent); icon != "" {
|
||||
return icon
|
||||
}
|
||||
|
||||
// Fallback: Try to get HuggingFace user avatar
|
||||
if model.Author != "" {
|
||||
if avatar := getHuggingFaceAvatarURL(model.Author); avatar != "" {
|
||||
return avatar
|
||||
}
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
2
.github/gallery-agent/gallery.go
vendored
2
.github/gallery-agent/gallery.go
vendored
@@ -7,8 +7,8 @@ import (
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/ghodss/yaml"
|
||||
"github.com/mudler/LocalAI/core/gallery/importers"
|
||||
"sigs.k8s.io/yaml"
|
||||
)
|
||||
|
||||
func formatTextContent(text string) string {
|
||||
|
||||
301
.github/gallery-agent/helpers.go
vendored
Normal file
301
.github/gallery-agent/helpers.go
vendored
Normal file
@@ -0,0 +1,301 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
hfapi "github.com/mudler/LocalAI/pkg/huggingface-api"
|
||||
"sigs.k8s.io/yaml"
|
||||
)
|
||||
|
||||
var galleryIndexPath = os.Getenv("GALLERY_INDEX_PATH")
|
||||
|
||||
// getGalleryIndexPath returns the gallery index file path, with a default fallback
|
||||
func getGalleryIndexPath() string {
|
||||
if galleryIndexPath != "" {
|
||||
return galleryIndexPath
|
||||
}
|
||||
return "gallery/index.yaml"
|
||||
}
|
||||
|
||||
type galleryModel struct {
|
||||
Name string `yaml:"name"`
|
||||
Urls []string `yaml:"urls"`
|
||||
}
|
||||
|
||||
// loadGalleryURLSet parses gallery/index.yaml once and returns the set of
|
||||
// HuggingFace model URLs already present in the gallery.
|
||||
func loadGalleryURLSet() (map[string]struct{}, error) {
|
||||
indexPath := getGalleryIndexPath()
|
||||
content, err := os.ReadFile(indexPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read %s: %w", indexPath, err)
|
||||
}
|
||||
|
||||
var galleryModels []galleryModel
|
||||
if err := yaml.Unmarshal(content, &galleryModels); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal %s: %w", indexPath, err)
|
||||
}
|
||||
|
||||
set := make(map[string]struct{}, len(galleryModels))
|
||||
for _, gm := range galleryModels {
|
||||
for _, u := range gm.Urls {
|
||||
set[u] = struct{}{}
|
||||
}
|
||||
}
|
||||
|
||||
// Also skip URLs already proposed in open (unmerged) gallery-agent PRs.
|
||||
// The workflow injects these via EXTRA_SKIP_URLS so we don't keep
|
||||
// re-proposing the same model every run while a PR is waiting to merge.
|
||||
for _, line := range strings.FieldsFunc(os.Getenv("EXTRA_SKIP_URLS"), func(r rune) bool {
|
||||
return r == '\n' || r == ',' || r == ' '
|
||||
}) {
|
||||
u := strings.TrimSpace(line)
|
||||
if u != "" {
|
||||
set[u] = struct{}{}
|
||||
}
|
||||
}
|
||||
|
||||
return set, nil
|
||||
}
|
||||
|
||||
// modelAlreadyInGallery checks whether a HuggingFace model repo is already
|
||||
// referenced in the gallery URL set.
|
||||
func modelAlreadyInGallery(set map[string]struct{}, modelID string) bool {
|
||||
_, ok := set["https://huggingface.co/"+modelID]
|
||||
return ok
|
||||
}
|
||||
|
||||
// baseModelFromTags returns the first `base_model:<repo>` value found in the
|
||||
// tag list, or "" if none is present. HuggingFace surfaces the base model
|
||||
// declared in the model card's YAML frontmatter as such a tag.
|
||||
func baseModelFromTags(tags []string) string {
|
||||
for _, t := range tags {
|
||||
if strings.HasPrefix(t, "base_model:") {
|
||||
return strings.TrimPrefix(t, "base_model:")
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// licenseFromTags returns the `license:<id>` value from the tag list, or "".
|
||||
func licenseFromTags(tags []string) string {
|
||||
for _, t := range tags {
|
||||
if strings.HasPrefix(t, "license:") {
|
||||
return strings.TrimPrefix(t, "license:")
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// curatedTags produces the gallery tag list from HuggingFace's raw tag set.
|
||||
// Always includes llm + gguf, then adds whitelisted family / capability
|
||||
// markers when they appear in the HF tag list.
|
||||
func curatedTags(hfTags []string) []string {
|
||||
whitelist := []string{
|
||||
"gpu", "cpu",
|
||||
"llama", "mistral", "mixtral", "qwen", "qwen2", "qwen3",
|
||||
"gemma", "gemma2", "gemma3", "phi", "phi3", "phi4",
|
||||
"deepseek", "yi", "falcon", "command-r",
|
||||
"vision", "multimodal", "code", "chat",
|
||||
"instruction-tuned", "reasoning", "thinking",
|
||||
}
|
||||
seen := map[string]struct{}{}
|
||||
out := []string{"llm", "gguf"}
|
||||
seen["llm"] = struct{}{}
|
||||
seen["gguf"] = struct{}{}
|
||||
|
||||
hfSet := map[string]struct{}{}
|
||||
for _, t := range hfTags {
|
||||
hfSet[strings.ToLower(t)] = struct{}{}
|
||||
}
|
||||
for _, w := range whitelist {
|
||||
if _, ok := hfSet[w]; ok {
|
||||
if _, dup := seen[w]; !dup {
|
||||
out = append(out, w)
|
||||
seen[w] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// resolveReadme fetches a description-quality README for a (possibly
|
||||
// quantized) repo: if a `base_model:` tag is present, fetch the base repo's
|
||||
// README; otherwise fall back to the repo's own README.
|
||||
func resolveReadme(client *hfapi.Client, modelID string, hfTags []string) (string, error) {
|
||||
if base := baseModelFromTags(hfTags); base != "" && base != modelID {
|
||||
if content, err := client.GetReadmeContent(base, "README.md"); err == nil && strings.TrimSpace(content) != "" {
|
||||
return cleanTextContent(content), nil
|
||||
}
|
||||
}
|
||||
content, err := client.GetReadmeContent(modelID, "README.md")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return cleanTextContent(content), nil
|
||||
}
|
||||
|
||||
// extractDescription turns a raw HuggingFace README into a concise plain-text
|
||||
// description suitable for embedding in gallery/index.yaml: strips YAML
|
||||
// frontmatter, HTML tags/comments, markdown images, link URLs (keeping the
|
||||
// link text), markdown tables, and then truncates at a paragraph boundary
|
||||
// around ~1200 characters. Raw README should still be used for icon
|
||||
// extraction — call this only for the `description:` field.
|
||||
func extractDescription(readme string) string {
|
||||
s := readme
|
||||
|
||||
// Strip leading YAML frontmatter: `---\n...\n---\n` at start of file.
|
||||
if strings.HasPrefix(strings.TrimLeft(s, " \t\n"), "---") {
|
||||
trimmed := strings.TrimLeft(s, " \t\n")
|
||||
rest := strings.TrimPrefix(trimmed, "---")
|
||||
if idx := strings.Index(rest, "\n---"); idx >= 0 {
|
||||
after := rest[idx+len("\n---"):]
|
||||
after = strings.TrimPrefix(after, "\n")
|
||||
s = after
|
||||
}
|
||||
}
|
||||
|
||||
// Strip HTML comments and tags.
|
||||
s = regexp.MustCompile(`(?s)<!--.*?-->`).ReplaceAllString(s, "")
|
||||
s = regexp.MustCompile(`(?is)<[^>]+>`).ReplaceAllString(s, "")
|
||||
|
||||
// Strip markdown images entirely.
|
||||
s = regexp.MustCompile(`!\[[^\]]*\]\([^)]*\)`).ReplaceAllString(s, "")
|
||||
// Replace markdown links `[text](url)` with just `text`.
|
||||
s = regexp.MustCompile(`\[([^\]]+)\]\([^)]+\)`).ReplaceAllString(s, "$1")
|
||||
|
||||
// Drop table lines and horizontal rules, and flatten all leading
|
||||
// whitespace: generateYAMLEntry embeds this under a `description: |`
|
||||
// literal block whose indentation is set by the first non-empty line.
|
||||
// If any line has extra leading whitespace (e.g. from an indented
|
||||
// `<p align="center">` block in the original README), YAML will pick
|
||||
// that up as the block's indent and every later line at a smaller
|
||||
// indent blows the block scalar. Stripping leading whitespace here
|
||||
// guarantees uniform 4-space indentation after formatTextContent runs.
|
||||
var kept []string
|
||||
for _, line := range strings.Split(s, "\n") {
|
||||
t := strings.TrimLeft(line, " \t")
|
||||
ts := strings.TrimSpace(t)
|
||||
if strings.HasPrefix(ts, "|") {
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(ts, ":--") || strings.HasPrefix(ts, "---") || strings.HasPrefix(ts, "===") {
|
||||
continue
|
||||
}
|
||||
kept = append(kept, t)
|
||||
}
|
||||
s = strings.Join(kept, "\n")
|
||||
|
||||
// Normalise whitespace and drop any leading blank lines so the literal
|
||||
// block in YAML doesn't start with a blank first line (which would
|
||||
// break the indentation detector the same way).
|
||||
s = cleanTextContent(s)
|
||||
s = strings.TrimLeft(s, " \t\n")
|
||||
|
||||
// Truncate at a paragraph boundary around maxLen chars.
|
||||
const maxLen = 1200
|
||||
if len(s) > maxLen {
|
||||
cut := strings.LastIndex(s[:maxLen], "\n\n")
|
||||
if cut < maxLen/3 {
|
||||
cut = maxLen
|
||||
}
|
||||
s = strings.TrimRight(s[:cut], " \t\n") + "\n\n..."
|
||||
}
|
||||
|
||||
return s
|
||||
}
|
||||
|
||||
// cleanTextContent removes trailing spaces/tabs and collapses multiple empty
|
||||
// lines so README content embeds cleanly into YAML without lint noise.
|
||||
func cleanTextContent(text string) string {
|
||||
lines := strings.Split(text, "\n")
|
||||
var cleaned []string
|
||||
var prevEmpty bool
|
||||
for _, line := range lines {
|
||||
trimmed := strings.TrimRight(line, " \t\r")
|
||||
if trimmed == "" {
|
||||
if !prevEmpty {
|
||||
cleaned = append(cleaned, "")
|
||||
}
|
||||
prevEmpty = true
|
||||
} else {
|
||||
cleaned = append(cleaned, trimmed)
|
||||
prevEmpty = false
|
||||
}
|
||||
}
|
||||
return strings.TrimRight(strings.Join(cleaned, "\n"), "\n")
|
||||
}
|
||||
|
||||
// extractIconFromReadme scans README content for an image URL usable as a
|
||||
// gallery entry icon.
|
||||
func extractIconFromReadme(readmeContent string) string {
|
||||
if readmeContent == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
markdownImageRegex := regexp.MustCompile(`(?i)!\[[^\]]*\]\(([^)]+\.(png|jpg|jpeg|svg|webp|gif))\)`)
|
||||
htmlImageRegex := regexp.MustCompile(`(?i)<img[^>]+src=["']([^"']+\.(png|jpg|jpeg|svg|webp|gif))["']`)
|
||||
plainImageRegex := regexp.MustCompile(`(?i)https?://[^\s<>"']+\.(png|jpg|jpeg|svg|webp|gif)`)
|
||||
|
||||
if m := markdownImageRegex.FindStringSubmatch(readmeContent); len(m) > 1 && strings.HasPrefix(strings.ToLower(m[1]), "http") {
|
||||
return strings.TrimSpace(m[1])
|
||||
}
|
||||
if m := htmlImageRegex.FindStringSubmatch(readmeContent); len(m) > 1 && strings.HasPrefix(strings.ToLower(m[1]), "http") {
|
||||
return strings.TrimSpace(m[1])
|
||||
}
|
||||
if m := plainImageRegex.FindStringSubmatch(readmeContent); len(m) > 0 && strings.HasPrefix(strings.ToLower(m[0]), "http") {
|
||||
return strings.TrimSpace(m[0])
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// getHuggingFaceAvatarURL returns the HF avatar URL for a user, or "".
|
||||
func getHuggingFaceAvatarURL(author string) string {
|
||||
if author == "" {
|
||||
return ""
|
||||
}
|
||||
userURL := fmt.Sprintf("https://huggingface.co/api/users/%s/overview", author)
|
||||
resp, err := http.Get(userURL)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return ""
|
||||
}
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
var info map[string]any
|
||||
if err := json.Unmarshal(body, &info); err != nil {
|
||||
return ""
|
||||
}
|
||||
if v, ok := info["avatarUrl"].(string); ok && v != "" {
|
||||
return v
|
||||
}
|
||||
if v, ok := info["avatar"].(string); ok && v != "" {
|
||||
return v
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// extractModelIcon extracts an icon URL from the README, falling back to the
|
||||
// HuggingFace user avatar.
|
||||
func extractModelIcon(model ProcessedModel) string {
|
||||
if icon := extractIconFromReadme(model.ReadmeContent); icon != "" {
|
||||
return icon
|
||||
}
|
||||
if model.Author != "" {
|
||||
if avatar := getHuggingFaceAvatarURL(model.Author); avatar != "" {
|
||||
return avatar
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
409
.github/gallery-agent/main.go
vendored
409
.github/gallery-agent/main.go
vendored
@@ -6,7 +6,6 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
hfapi "github.com/mudler/LocalAI/pkg/huggingface-api"
|
||||
@@ -39,16 +38,6 @@ type ProcessedModel struct {
|
||||
Icon string `json:"icon,omitempty"`
|
||||
}
|
||||
|
||||
// SearchResult represents the complete result of searching and processing models
|
||||
type SearchResult struct {
|
||||
SearchTerm string `json:"search_term"`
|
||||
Limit int `json:"limit"`
|
||||
Quantization string `json:"quantization"`
|
||||
TotalModelsFound int `json:"total_models_found"`
|
||||
Models []ProcessedModel `json:"models"`
|
||||
FormattedOutput string `json:"formatted_output"`
|
||||
}
|
||||
|
||||
// AddedModelSummary represents a summary of models added to the gallery
|
||||
type AddedModelSummary struct {
|
||||
SearchTerm string `json:"search_term"`
|
||||
@@ -63,19 +52,16 @@ type AddedModelSummary struct {
|
||||
func main() {
|
||||
startTime := time.Now()
|
||||
|
||||
// Check for synthetic mode
|
||||
syntheticMode := os.Getenv("SYNTHETIC_MODE")
|
||||
if syntheticMode == "true" || syntheticMode == "1" {
|
||||
// Synthetic mode for local testing
|
||||
if sm := os.Getenv("SYNTHETIC_MODE"); sm == "true" || sm == "1" {
|
||||
fmt.Println("Running in SYNTHETIC MODE - generating random test data")
|
||||
err := runSyntheticMode()
|
||||
if err != nil {
|
||||
if err := runSyntheticMode(); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error in synthetic mode: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// Get configuration from environment variables
|
||||
searchTerm := os.Getenv("SEARCH_TERM")
|
||||
if searchTerm == "" {
|
||||
searchTerm = "GGUF"
|
||||
@@ -83,7 +69,7 @@ func main() {
|
||||
|
||||
limitStr := os.Getenv("LIMIT")
|
||||
if limitStr == "" {
|
||||
limitStr = "5"
|
||||
limitStr = "15"
|
||||
}
|
||||
limit, err := strconv.Atoi(limitStr)
|
||||
if err != nil {
|
||||
@@ -92,287 +78,197 @@ func main() {
|
||||
}
|
||||
|
||||
quantization := os.Getenv("QUANTIZATION")
|
||||
|
||||
maxModels := os.Getenv("MAX_MODELS")
|
||||
if maxModels == "" {
|
||||
maxModels = "1"
|
||||
if quantization == "" {
|
||||
quantization = "Q4_K_M"
|
||||
}
|
||||
maxModelsInt, err := strconv.Atoi(maxModels)
|
||||
|
||||
maxModelsStr := os.Getenv("MAX_MODELS")
|
||||
if maxModelsStr == "" {
|
||||
maxModelsStr = "1"
|
||||
}
|
||||
maxModels, err := strconv.Atoi(maxModelsStr)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error parsing MAX_MODELS: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Print configuration
|
||||
fmt.Printf("Gallery Agent Configuration:\n")
|
||||
fmt.Printf(" Search Term: %s\n", searchTerm)
|
||||
fmt.Printf(" Limit: %d\n", limit)
|
||||
fmt.Printf(" Quantization: %s\n", quantization)
|
||||
fmt.Printf(" Max Models to Add: %d\n", maxModelsInt)
|
||||
fmt.Printf(" Gallery Index Path: %s\n", os.Getenv("GALLERY_INDEX_PATH"))
|
||||
fmt.Printf(" Max Models to Add: %d\n", maxModels)
|
||||
fmt.Printf(" Gallery Index Path: %s\n", getGalleryIndexPath())
|
||||
fmt.Println()
|
||||
|
||||
result, err := searchAndProcessModels(searchTerm, limit, quantization)
|
||||
// Phase 1: load current gallery and query HuggingFace.
|
||||
gallerySet, err := loadGalleryURLSet()
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
|
||||
fmt.Fprintf(os.Stderr, "Error loading gallery index: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Printf("Loaded %d existing gallery entries\n", len(gallerySet))
|
||||
|
||||
fmt.Println(result.FormattedOutput)
|
||||
var models []ProcessedModel
|
||||
|
||||
if len(result.Models) > 1 {
|
||||
fmt.Println("More than one model found (", len(result.Models), "), using AI agent to select the most interesting models")
|
||||
for _, model := range result.Models {
|
||||
fmt.Println("Model: ", model.ModelID)
|
||||
}
|
||||
// Use AI agent to select the most interesting models
|
||||
fmt.Println("Using AI agent to select the most interesting models...")
|
||||
models, err = selectMostInterestingModels(context.Background(), result)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error in model selection: %v\n", err)
|
||||
// Continue with original result if selection fails
|
||||
models = result.Models
|
||||
}
|
||||
} else if len(result.Models) == 1 {
|
||||
models = result.Models
|
||||
fmt.Println("Only one model found, using it directly")
|
||||
}
|
||||
|
||||
fmt.Print(models)
|
||||
|
||||
// Filter out models that already exist in the gallery
|
||||
fmt.Println("Filtering out existing models...")
|
||||
models, err = filterExistingModels(models)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error filtering existing models: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Limit to maxModelsInt after filtering
|
||||
if len(models) > maxModelsInt {
|
||||
models = models[:maxModelsInt]
|
||||
}
|
||||
|
||||
// Track added models for summary
|
||||
var addedModelIDs []string
|
||||
var addedModelURLs []string
|
||||
|
||||
// Generate YAML entries and append to gallery/index.yaml
|
||||
if len(models) > 0 {
|
||||
for _, model := range models {
|
||||
addedModelIDs = append(addedModelIDs, model.ModelID)
|
||||
// Generate Hugging Face URL for the model
|
||||
modelURL := fmt.Sprintf("https://huggingface.co/%s", model.ModelID)
|
||||
addedModelURLs = append(addedModelURLs, modelURL)
|
||||
}
|
||||
fmt.Println("Generating YAML entries for selected models...")
|
||||
err = generateYAMLForModels(context.Background(), models, quantization)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error generating YAML entries: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
} else {
|
||||
fmt.Println("No new models to add to the gallery.")
|
||||
}
|
||||
|
||||
// Create and write summary
|
||||
processingTime := time.Since(startTime).String()
|
||||
summary := AddedModelSummary{
|
||||
SearchTerm: searchTerm,
|
||||
TotalFound: result.TotalModelsFound,
|
||||
ModelsAdded: len(addedModelIDs),
|
||||
AddedModelIDs: addedModelIDs,
|
||||
AddedModelURLs: addedModelURLs,
|
||||
Quantization: quantization,
|
||||
ProcessingTime: processingTime,
|
||||
}
|
||||
|
||||
// Write summary to file
|
||||
summaryData, err := json.MarshalIndent(summary, "", " ")
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error marshaling summary: %v\n", err)
|
||||
} else {
|
||||
err = os.WriteFile("gallery-agent-summary.json", summaryData, 0644)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error writing summary file: %v\n", err)
|
||||
} else {
|
||||
fmt.Printf("Summary written to gallery-agent-summary.json\n")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func searchAndProcessModels(searchTerm string, limit int, quantization string) (*SearchResult, error) {
|
||||
client := hfapi.NewClient()
|
||||
var outputBuilder strings.Builder
|
||||
|
||||
fmt.Println("Searching for models...")
|
||||
// Initialize the result struct
|
||||
result := &SearchResult{
|
||||
SearchTerm: searchTerm,
|
||||
Limit: limit,
|
||||
Quantization: quantization,
|
||||
Models: []ProcessedModel{},
|
||||
}
|
||||
|
||||
models, err := client.GetLatest(searchTerm, limit)
|
||||
fmt.Println("Searching for trending models on HuggingFace...")
|
||||
rawModels, err := client.GetTrending(searchTerm, limit)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to fetch models: %w", err)
|
||||
fmt.Fprintf(os.Stderr, "Error fetching models: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Printf("Found %d trending models matching %q\n", len(rawModels), searchTerm)
|
||||
totalFound := len(rawModels)
|
||||
|
||||
// Phase 2: drop anything already in the gallery *before* any expensive
|
||||
// per-model work (GetModelDetails, README fetches, icon lookups).
|
||||
fresh := rawModels[:0]
|
||||
for _, m := range rawModels {
|
||||
if modelAlreadyInGallery(gallerySet, m.ModelID) {
|
||||
fmt.Printf("Skipping existing model: %s\n", m.ModelID)
|
||||
continue
|
||||
}
|
||||
fresh = append(fresh, m)
|
||||
}
|
||||
fmt.Printf("%d candidates after gallery dedup\n", len(fresh))
|
||||
|
||||
// Phase 3: HuggingFace already returned these in trendingScore order —
|
||||
// just cap to MAX_MODELS.
|
||||
if len(fresh) > maxModels {
|
||||
fresh = fresh[:maxModels]
|
||||
}
|
||||
if len(fresh) == 0 {
|
||||
fmt.Println("No new models to add to the gallery.")
|
||||
writeSummary(AddedModelSummary{
|
||||
SearchTerm: searchTerm,
|
||||
TotalFound: totalFound,
|
||||
ModelsAdded: 0,
|
||||
Quantization: quantization,
|
||||
ProcessingTime: time.Since(startTime).String(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
fmt.Println("Models found:", len(models))
|
||||
result.TotalModelsFound = len(models)
|
||||
// Phase 4: fetch details and build ProcessedModel entries for survivors.
|
||||
var processed []ProcessedModel
|
||||
quantPrefs := []string{quantization, "Q4_K_M", "Q4_K_S", "Q3_K_M", "Q2_K", "Q8_0"}
|
||||
for _, m := range fresh {
|
||||
fmt.Printf("Processing model: %s (downloads=%d)\n", m.ModelID, m.Downloads)
|
||||
|
||||
if len(models) == 0 {
|
||||
outputBuilder.WriteString("No models found.\n")
|
||||
result.FormattedOutput = outputBuilder.String()
|
||||
return result, nil
|
||||
}
|
||||
|
||||
outputBuilder.WriteString(fmt.Sprintf("Found %d models matching '%s':\n\n", len(models), searchTerm))
|
||||
|
||||
// Process each model
|
||||
for i, model := range models {
|
||||
outputBuilder.WriteString(fmt.Sprintf("%d. Processing Model: %s\n", i+1, model.ModelID))
|
||||
outputBuilder.WriteString(fmt.Sprintf(" Author: %s\n", model.Author))
|
||||
outputBuilder.WriteString(fmt.Sprintf(" Downloads: %d\n", model.Downloads))
|
||||
outputBuilder.WriteString(fmt.Sprintf(" Last Modified: %s\n", model.LastModified))
|
||||
|
||||
// Initialize processed model struct
|
||||
processedModel := ProcessedModel{
|
||||
ModelID: model.ModelID,
|
||||
Author: model.Author,
|
||||
Downloads: model.Downloads,
|
||||
LastModified: model.LastModified,
|
||||
QuantizationPreferences: []string{quantization, "Q4_K_M", "Q4_K_S", "Q3_K_M", "Q2_K"},
|
||||
pm := ProcessedModel{
|
||||
ModelID: m.ModelID,
|
||||
Author: m.Author,
|
||||
Downloads: m.Downloads,
|
||||
LastModified: m.LastModified,
|
||||
QuantizationPreferences: quantPrefs,
|
||||
}
|
||||
|
||||
// Get detailed model information
|
||||
details, err := client.GetModelDetails(model.ModelID)
|
||||
details, err := client.GetModelDetails(m.ModelID)
|
||||
if err != nil {
|
||||
errorMsg := fmt.Sprintf(" Error getting model details: %v\n", err)
|
||||
outputBuilder.WriteString(errorMsg)
|
||||
processedModel.ProcessingError = err.Error()
|
||||
result.Models = append(result.Models, processedModel)
|
||||
fmt.Printf(" Error getting model details: %v (skipping)\n", err)
|
||||
continue
|
||||
}
|
||||
|
||||
// Define quantization preferences (in order of preference)
|
||||
quantizationPreferences := []string{quantization, "Q4_K_M", "Q4_K_S", "Q3_K_M", "Q2_K"}
|
||||
preferred := hfapi.FindPreferredModelFile(details.Files, quantPrefs)
|
||||
if preferred == nil {
|
||||
fmt.Printf(" No GGUF file matching %v — skipping\n", quantPrefs)
|
||||
continue
|
||||
}
|
||||
|
||||
// Find preferred model file
|
||||
preferredModelFile := hfapi.FindPreferredModelFile(details.Files, quantizationPreferences)
|
||||
|
||||
// Process files
|
||||
processedFiles := make([]ProcessedModelFile, len(details.Files))
|
||||
for j, file := range details.Files {
|
||||
pm.Files = make([]ProcessedModelFile, len(details.Files))
|
||||
for j, f := range details.Files {
|
||||
fileType := "other"
|
||||
if file.IsReadme {
|
||||
if f.IsReadme {
|
||||
fileType = "readme"
|
||||
} else if preferredModelFile != nil && file.Path == preferredModelFile.Path {
|
||||
} else if f.Path == preferred.Path {
|
||||
fileType = "model"
|
||||
}
|
||||
|
||||
processedFiles[j] = ProcessedModelFile{
|
||||
Path: file.Path,
|
||||
Size: file.Size,
|
||||
SHA256: file.SHA256,
|
||||
IsReadme: file.IsReadme,
|
||||
pm.Files[j] = ProcessedModelFile{
|
||||
Path: f.Path,
|
||||
Size: f.Size,
|
||||
SHA256: f.SHA256,
|
||||
IsReadme: f.IsReadme,
|
||||
FileType: fileType,
|
||||
}
|
||||
}
|
||||
|
||||
processedModel.Files = processedFiles
|
||||
|
||||
// Set preferred model file
|
||||
if preferredModelFile != nil {
|
||||
for _, file := range processedFiles {
|
||||
if file.Path == preferredModelFile.Path {
|
||||
processedModel.PreferredModelFile = &file
|
||||
break
|
||||
}
|
||||
if f.Path == preferred.Path {
|
||||
copyFile := pm.Files[j]
|
||||
pm.PreferredModelFile = ©File
|
||||
}
|
||||
if f.IsReadme {
|
||||
copyFile := pm.Files[j]
|
||||
pm.ReadmeFile = ©File
|
||||
}
|
||||
}
|
||||
|
||||
// Print file information
|
||||
outputBuilder.WriteString(fmt.Sprintf(" Files found: %d\n", len(details.Files)))
|
||||
// Deterministic README resolution: follow base_model tag if set.
|
||||
// Keep the raw (HTML-bearing) README around while we extract the
|
||||
// icon, then strip it down to a plain-text description for the
|
||||
// `description:` YAML field.
|
||||
readme, err := resolveReadme(client, m.ModelID, m.Tags)
|
||||
if err != nil {
|
||||
fmt.Printf(" Warning: failed to fetch README: %v\n", err)
|
||||
}
|
||||
pm.ReadmeContent = readme
|
||||
|
||||
if preferredModelFile != nil {
|
||||
outputBuilder.WriteString(fmt.Sprintf(" Preferred Model File: %s (SHA256: %s)\n",
|
||||
preferredModelFile.Path,
|
||||
preferredModelFile.SHA256))
|
||||
} else {
|
||||
outputBuilder.WriteString(fmt.Sprintf(" No model file found with quantization preferences: %v\n", quantizationPreferences))
|
||||
pm.License = licenseFromTags(m.Tags)
|
||||
pm.Tags = curatedTags(m.Tags)
|
||||
pm.Icon = extractModelIcon(pm)
|
||||
|
||||
if pm.ReadmeContent != "" {
|
||||
pm.ReadmeContent = extractDescription(pm.ReadmeContent)
|
||||
pm.ReadmeContentPreview = truncateString(pm.ReadmeContent, 200)
|
||||
}
|
||||
|
||||
if details.ReadmeFile != nil {
|
||||
outputBuilder.WriteString(fmt.Sprintf(" README File: %s\n", details.ReadmeFile.Path))
|
||||
|
||||
// Find and set readme file
|
||||
for _, file := range processedFiles {
|
||||
if file.IsReadme {
|
||||
processedModel.ReadmeFile = &file
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Println("Getting real readme for", model.ModelID, "waiting...")
|
||||
// Use agent to get the real readme and prepare the model description
|
||||
readmeContent, err := getRealReadme(context.Background(), model.ModelID)
|
||||
if err == nil {
|
||||
processedModel.ReadmeContent = readmeContent
|
||||
processedModel.ReadmeContentPreview = truncateString(readmeContent, 200)
|
||||
outputBuilder.WriteString(fmt.Sprintf(" README Content Preview: %s\n",
|
||||
processedModel.ReadmeContentPreview))
|
||||
} else {
|
||||
fmt.Printf(" Warning: Failed to get real readme: %v\n", err)
|
||||
}
|
||||
fmt.Println("Real readme got", readmeContent)
|
||||
|
||||
// Extract metadata (tags, license) from README using LLM
|
||||
fmt.Println("Extracting metadata for", model.ModelID, "waiting...")
|
||||
tags, license, err := extractModelMetadata(context.Background(), processedModel)
|
||||
if err == nil {
|
||||
processedModel.Tags = tags
|
||||
processedModel.License = license
|
||||
outputBuilder.WriteString(fmt.Sprintf(" Tags: %v\n", tags))
|
||||
outputBuilder.WriteString(fmt.Sprintf(" License: %s\n", license))
|
||||
} else {
|
||||
fmt.Printf(" Warning: Failed to extract metadata: %v\n", err)
|
||||
}
|
||||
|
||||
// Extract icon from README or use HuggingFace avatar
|
||||
icon := extractModelIcon(processedModel)
|
||||
if icon != "" {
|
||||
processedModel.Icon = icon
|
||||
outputBuilder.WriteString(fmt.Sprintf(" Icon: %s\n", icon))
|
||||
}
|
||||
// Get README content
|
||||
// readmeContent, err := client.GetReadmeContent(model.ModelID, details.ReadmeFile.Path)
|
||||
// if err == nil {
|
||||
// processedModel.ReadmeContent = readmeContent
|
||||
// processedModel.ReadmeContentPreview = truncateString(readmeContent, 200)
|
||||
// outputBuilder.WriteString(fmt.Sprintf(" README Content Preview: %s\n",
|
||||
// processedModel.ReadmeContentPreview))
|
||||
// }
|
||||
}
|
||||
|
||||
// Print all files with their checksums
|
||||
outputBuilder.WriteString(" All Files:\n")
|
||||
for _, file := range processedFiles {
|
||||
outputBuilder.WriteString(fmt.Sprintf(" - %s (%s, %d bytes", file.Path, file.FileType, file.Size))
|
||||
if file.SHA256 != "" {
|
||||
outputBuilder.WriteString(fmt.Sprintf(", SHA256: %s", file.SHA256))
|
||||
}
|
||||
outputBuilder.WriteString(")\n")
|
||||
}
|
||||
|
||||
outputBuilder.WriteString("\n")
|
||||
result.Models = append(result.Models, processedModel)
|
||||
fmt.Printf(" License: %s, Tags: %v, Icon: %s\n", pm.License, pm.Tags, pm.Icon)
|
||||
processed = append(processed, pm)
|
||||
}
|
||||
|
||||
result.FormattedOutput = outputBuilder.String()
|
||||
return result, nil
|
||||
if len(processed) == 0 {
|
||||
fmt.Println("No processable models after detail fetch.")
|
||||
writeSummary(AddedModelSummary{
|
||||
SearchTerm: searchTerm,
|
||||
TotalFound: totalFound,
|
||||
ModelsAdded: 0,
|
||||
Quantization: quantization,
|
||||
ProcessingTime: time.Since(startTime).String(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Phase 5: write YAML entries.
|
||||
var addedIDs, addedURLs []string
|
||||
for _, pm := range processed {
|
||||
addedIDs = append(addedIDs, pm.ModelID)
|
||||
addedURLs = append(addedURLs, "https://huggingface.co/"+pm.ModelID)
|
||||
}
|
||||
|
||||
fmt.Println("Generating YAML entries for selected models...")
|
||||
if err := generateYAMLForModels(context.Background(), processed, quantization); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error generating YAML entries: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
writeSummary(AddedModelSummary{
|
||||
SearchTerm: searchTerm,
|
||||
TotalFound: totalFound,
|
||||
ModelsAdded: len(addedIDs),
|
||||
AddedModelIDs: addedIDs,
|
||||
AddedModelURLs: addedURLs,
|
||||
Quantization: quantization,
|
||||
ProcessingTime: time.Since(startTime).String(),
|
||||
})
|
||||
}
|
||||
|
||||
func writeSummary(summary AddedModelSummary) {
|
||||
data, err := json.MarshalIndent(summary, "", " ")
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error marshaling summary: %v\n", err)
|
||||
return
|
||||
}
|
||||
if err := os.WriteFile("gallery-agent-summary.json", data, 0644); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error writing summary file: %v\n", err)
|
||||
return
|
||||
}
|
||||
fmt.Println("Summary written to gallery-agent-summary.json")
|
||||
}
|
||||
|
||||
func truncateString(s string, maxLen int) string {
|
||||
@@ -381,3 +277,4 @@ func truncateString(s string, maxLen int) string {
|
||||
}
|
||||
return s[:maxLen] + "..."
|
||||
}
|
||||
|
||||
|
||||
46
.github/gallery-agent/tools.go
vendored
46
.github/gallery-agent/tools.go
vendored
@@ -1,46 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
hfapi "github.com/mudler/LocalAI/pkg/huggingface-api"
|
||||
openai "github.com/sashabaranov/go-openai"
|
||||
jsonschema "github.com/sashabaranov/go-openai/jsonschema"
|
||||
)
|
||||
|
||||
// Get repository README from HF
|
||||
type HFReadmeTool struct {
|
||||
client *hfapi.Client
|
||||
}
|
||||
|
||||
func (s *HFReadmeTool) Execute(args map[string]any) (string, any, error) {
|
||||
q, ok := args["repository"].(string)
|
||||
if !ok {
|
||||
return "", nil, fmt.Errorf("no query")
|
||||
}
|
||||
readme, err := s.client.GetReadmeContent(q, "README.md")
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
return readme, nil, nil
|
||||
}
|
||||
|
||||
func (s *HFReadmeTool) Tool() openai.Tool {
|
||||
return openai.Tool{
|
||||
Type: openai.ToolTypeFunction,
|
||||
Function: &openai.FunctionDefinition{
|
||||
Name: "hf_readme",
|
||||
Description: "A tool to get the README content of a huggingface repository",
|
||||
Parameters: jsonschema.Definition{
|
||||
Type: jsonschema.Object,
|
||||
Properties: map[string]jsonschema.Definition{
|
||||
"repository": {
|
||||
Type: jsonschema.String,
|
||||
Description: "The huggingface repository to get the README content of",
|
||||
},
|
||||
},
|
||||
Required: []string{"repository"},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
430
.github/workflows/backend.yml
vendored
430
.github/workflows/backend.yml
vendored
@@ -30,6 +30,7 @@ jobs:
|
||||
skip-drivers: ${{ matrix.skip-drivers }}
|
||||
context: ${{ matrix.context }}
|
||||
ubuntu-version: ${{ matrix.ubuntu-version }}
|
||||
amdgpu-targets: ${{ matrix.amdgpu-targets }}
|
||||
secrets:
|
||||
dockerUsername: ${{ secrets.DOCKERHUB_USERNAME }}
|
||||
dockerPassword: ${{ secrets.DOCKERHUB_PASSWORD }}
|
||||
@@ -66,6 +67,19 @@ jobs:
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-cpu-sglang'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'true'
|
||||
backend: "sglang"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
@@ -105,6 +119,25 @@ jobs:
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
# tinygrad ships a single image — its CPU device uses bundled
|
||||
# libLLVM, and its CUDA / HIP / Metal devices dlopen the host
|
||||
# driver libraries at runtime via tinygrad's ctypes autogen
|
||||
# wrappers. There is no toolkit-version split because tinygrad
|
||||
# generates kernels itself (PTX renderer for CUDA) and never
|
||||
# links against cuDNN/cuBLAS/torch.
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-tinygrad'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'true'
|
||||
backend: "tinygrad"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
@@ -353,6 +386,32 @@ jobs:
|
||||
dockerfile: "./backend/Dockerfile.llama-cpp"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "8"
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-nvidia-cuda-12-turboquant'
|
||||
runs-on: 'bigger-runner'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "turboquant"
|
||||
dockerfile: "./backend/Dockerfile.turboquant"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "8"
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-nvidia-cuda-12-buun-llama-cpp'
|
||||
runs-on: 'bigger-runner'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "buun-llama-cpp"
|
||||
dockerfile: "./backend/Dockerfile.buun-llama-cpp"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "8"
|
||||
@@ -379,6 +438,19 @@ jobs:
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "8"
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-nvidia-cuda-12-sglang'
|
||||
runs-on: 'arc-runner-set'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "sglang"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "8"
|
||||
@@ -652,6 +724,32 @@ jobs:
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "8"
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-nvidia-cuda-12-insightface'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "insightface"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "8"
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-nvidia-cuda-12-speaker-recognition'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "speaker-recognition"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "8"
|
||||
@@ -796,6 +894,32 @@ jobs:
|
||||
dockerfile: "./backend/Dockerfile.llama-cpp"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-nvidia-cuda-13-turboquant'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "turboquant"
|
||||
dockerfile: "./backend/Dockerfile.turboquant"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-nvidia-cuda-13-buun-llama-cpp'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "buun-llama-cpp"
|
||||
dockerfile: "./backend/Dockerfile.buun-llama-cpp"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
@@ -809,6 +933,32 @@ jobs:
|
||||
backend: "llama-cpp"
|
||||
dockerfile: "./backend/Dockerfile.llama-cpp"
|
||||
context: "./"
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/arm64'
|
||||
skip-drivers: 'false'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-nvidia-l4t-cuda-13-arm64-turboquant'
|
||||
base-image: "ubuntu:24.04"
|
||||
runs-on: 'ubuntu-24.04-arm'
|
||||
ubuntu-version: '2404'
|
||||
backend: "turboquant"
|
||||
dockerfile: "./backend/Dockerfile.turboquant"
|
||||
context: "./"
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/arm64'
|
||||
skip-drivers: 'false'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-nvidia-l4t-cuda-13-arm64-buun-llama-cpp'
|
||||
base-image: "ubuntu:24.04"
|
||||
runs-on: 'ubuntu-24.04-arm'
|
||||
ubuntu-version: '2404'
|
||||
backend: "buun-llama-cpp"
|
||||
dockerfile: "./backend/Dockerfile.buun-llama-cpp"
|
||||
context: "./"
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
@@ -1330,6 +1480,32 @@ jobs:
|
||||
dockerfile: "./backend/Dockerfile.llama-cpp"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'hipblas'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-rocm-hipblas-turboquant'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "rocm/dev-ubuntu-24.04:7.2.1"
|
||||
skip-drivers: 'false'
|
||||
backend: "turboquant"
|
||||
dockerfile: "./backend/Dockerfile.turboquant"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'hipblas'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-rocm-hipblas-buun-llama-cpp'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "rocm/dev-ubuntu-24.04:7.2.1"
|
||||
skip-drivers: 'false'
|
||||
backend: "buun-llama-cpp"
|
||||
dockerfile: "./backend/Dockerfile.buun-llama-cpp"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'hipblas'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
@@ -1356,6 +1532,19 @@ jobs:
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'hipblas'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-rocm-hipblas-sglang'
|
||||
runs-on: 'arc-runner-set'
|
||||
base-image: "rocm/dev-ubuntu-24.04:7.2.1"
|
||||
skip-drivers: 'false'
|
||||
backend: "sglang"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'hipblas'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
@@ -1513,19 +1702,6 @@ jobs:
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'hipblas'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-rocm-hipblas-whisperx'
|
||||
runs-on: 'bigger-runner'
|
||||
base-image: "rocm/dev-ubuntu-24.04:7.2.1"
|
||||
skip-drivers: 'false'
|
||||
backend: "whisperx"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'hipblas'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
@@ -1566,6 +1742,32 @@ jobs:
|
||||
dockerfile: "./backend/Dockerfile.llama-cpp"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'sycl_f32'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-intel-sycl-f32-turboquant'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "intel/oneapi-basekit:2025.3.0-0-devel-ubuntu24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "turboquant"
|
||||
dockerfile: "./backend/Dockerfile.turboquant"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'sycl_f32'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-intel-sycl-f32-buun-llama-cpp'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "intel/oneapi-basekit:2025.3.0-0-devel-ubuntu24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "buun-llama-cpp"
|
||||
dockerfile: "./backend/Dockerfile.buun-llama-cpp"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'sycl_f16'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
@@ -1579,6 +1781,32 @@ jobs:
|
||||
dockerfile: "./backend/Dockerfile.llama-cpp"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'sycl_f16'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-intel-sycl-f16-turboquant'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "intel/oneapi-basekit:2025.3.0-0-devel-ubuntu24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "turboquant"
|
||||
dockerfile: "./backend/Dockerfile.turboquant"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'sycl_f16'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-intel-sycl-f16-buun-llama-cpp'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "intel/oneapi-basekit:2025.3.0-0-devel-ubuntu24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "buun-llama-cpp"
|
||||
dockerfile: "./backend/Dockerfile.buun-llama-cpp"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'intel'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
@@ -1592,6 +1820,19 @@ jobs:
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'intel'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-intel-sglang'
|
||||
runs-on: 'arc-runner-set'
|
||||
base-image: "intel/oneapi-basekit:2025.3.0-0-devel-ubuntu24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "sglang"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'intel'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
@@ -1958,6 +2199,32 @@ jobs:
|
||||
dockerfile: "./backend/Dockerfile.llama-cpp"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64,linux/arm64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-cpu-turboquant'
|
||||
runs-on: 'bigger-runner'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "turboquant"
|
||||
dockerfile: "./backend/Dockerfile.turboquant"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64,linux/arm64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-cpu-buun-llama-cpp'
|
||||
runs-on: 'bigger-runner'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "buun-llama-cpp"
|
||||
dockerfile: "./backend/Dockerfile.buun-llama-cpp"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
@@ -1984,6 +2251,32 @@ jobs:
|
||||
dockerfile: "./backend/Dockerfile.llama-cpp"
|
||||
context: "./"
|
||||
ubuntu-version: '2204'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/arm64'
|
||||
skip-drivers: 'false'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-nvidia-l4t-arm64-turboquant'
|
||||
base-image: "nvcr.io/nvidia/l4t-jetpack:r36.4.0"
|
||||
runs-on: 'ubuntu-24.04-arm'
|
||||
backend: "turboquant"
|
||||
dockerfile: "./backend/Dockerfile.turboquant"
|
||||
context: "./"
|
||||
ubuntu-version: '2204'
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/arm64'
|
||||
skip-drivers: 'false'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-nvidia-l4t-arm64-buun-llama-cpp'
|
||||
base-image: "nvcr.io/nvidia/l4t-jetpack:r36.4.0"
|
||||
runs-on: 'ubuntu-24.04-arm'
|
||||
backend: "buun-llama-cpp"
|
||||
dockerfile: "./backend/Dockerfile.buun-llama-cpp"
|
||||
context: "./"
|
||||
ubuntu-version: '2204'
|
||||
- build-type: 'vulkan'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
@@ -1997,6 +2290,32 @@ jobs:
|
||||
dockerfile: "./backend/Dockerfile.llama-cpp"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'vulkan'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64,linux/arm64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-vulkan-turboquant'
|
||||
runs-on: 'bigger-runner'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "turboquant"
|
||||
dockerfile: "./backend/Dockerfile.turboquant"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'vulkan'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64,linux/arm64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-vulkan-buun-llama-cpp'
|
||||
runs-on: 'bigger-runner'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "buun-llama-cpp"
|
||||
dockerfile: "./backend/Dockerfile.buun-llama-cpp"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
# Stablediffusion-ggml
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
@@ -2408,6 +2727,20 @@ jobs:
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
# kokoros (Rust TTS)
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-cpu-kokoros'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "kokoros"
|
||||
dockerfile: "./backend/Dockerfile.rust"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
# local-store
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
@@ -2436,6 +2769,34 @@ jobs:
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
# insightface (face recognition)
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64,linux/arm64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-cpu-insightface'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "insightface"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
# speaker-recognition (voice/speaker biometrics)
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64,linux/arm64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-cpu-speaker-recognition'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "speaker-recognition"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
- build-type: 'intel'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
@@ -2633,6 +2994,49 @@ jobs:
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
# sherpa-onnx CPU
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64,linux/arm64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-cpu-sherpa-onnx'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "sherpa-onnx"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
# sherpa-onnx CUDA 12
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "8"
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-nvidia-cuda-12-sherpa-onnx'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "sherpa-onnx"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
# sherpa-onnx CUDA 13 — requires onnxruntime 1.24.x+ for the
|
||||
# gpu_cuda13 tarball; sherpa-onnx SHERPA_COMMIT pins to v1.12.39.
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "13"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-nvidia-cuda-13-sherpa-onnx'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:24.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "sherpa-onnx"
|
||||
dockerfile: "./backend/Dockerfile.golang"
|
||||
context: "./"
|
||||
ubuntu-version: '2404'
|
||||
backend-jobs-darwin:
|
||||
uses: ./.github/workflows/backend_build_darwin.yml
|
||||
strategy:
|
||||
|
||||
9
.github/workflows/backend_build.yml
vendored
9
.github/workflows/backend_build.yml
vendored
@@ -58,6 +58,11 @@ on:
|
||||
required: false
|
||||
default: '2204'
|
||||
type: string
|
||||
amdgpu-targets:
|
||||
description: 'AMD GPU targets for ROCm/HIP builds'
|
||||
required: false
|
||||
default: 'gfx908,gfx90a,gfx942,gfx950,gfx1030,gfx1100,gfx1101,gfx1102,gfx1151,gfx1200,gfx1201'
|
||||
type: string
|
||||
secrets:
|
||||
dockerUsername:
|
||||
required: false
|
||||
@@ -103,6 +108,8 @@ jobs:
|
||||
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
submodules: true
|
||||
|
||||
- name: Release space from worker
|
||||
if: inputs.runs-on == 'ubuntu-latest'
|
||||
@@ -214,6 +221,7 @@ jobs:
|
||||
BASE_IMAGE=${{ inputs.base-image }}
|
||||
BACKEND=${{ inputs.backend }}
|
||||
UBUNTU_VERSION=${{ inputs.ubuntu-version }}
|
||||
AMDGPU_TARGETS=${{ inputs.amdgpu-targets }}
|
||||
context: ${{ inputs.context }}
|
||||
file: ${{ inputs.dockerfile }}
|
||||
cache-from: type=gha
|
||||
@@ -235,6 +243,7 @@ jobs:
|
||||
BASE_IMAGE=${{ inputs.base-image }}
|
||||
BACKEND=${{ inputs.backend }}
|
||||
UBUNTU_VERSION=${{ inputs.ubuntu-version }}
|
||||
AMDGPU_TARGETS=${{ inputs.amdgpu-targets }}
|
||||
context: ${{ inputs.context }}
|
||||
file: ${{ inputs.dockerfile }}
|
||||
cache-from: type=gha
|
||||
|
||||
4
.github/workflows/bump_deps.yaml
vendored
4
.github/workflows/bump_deps.yaml
vendored
@@ -18,6 +18,10 @@ jobs:
|
||||
variable: "IK_LLAMA_VERSION"
|
||||
branch: "main"
|
||||
file: "backend/cpp/ik-llama-cpp/Makefile"
|
||||
- repository: "TheTom/llama-cpp-turboquant"
|
||||
variable: "TURBOQUANT_VERSION"
|
||||
branch: "feature/turboquant-kv-cache"
|
||||
file: "backend/cpp/turboquant/Makefile"
|
||||
- repository: "ggml-org/whisper.cpp"
|
||||
variable: "WHISPER_CPP_VERSION"
|
||||
branch: "master"
|
||||
|
||||
99
.github/workflows/gallery-agent.yaml
vendored
99
.github/workflows/gallery-agent.yaml
vendored
@@ -48,21 +48,88 @@ jobs:
|
||||
go install google.golang.org/protobuf/cmd/protoc-gen-go@v1.34.2
|
||||
go install google.golang.org/grpc/cmd/protoc-gen-go-grpc@1958fcbe2ca8bd93af633f11e97d44e567e945af
|
||||
PATH="$PATH:$HOME/go/bin" make protogen-go
|
||||
- uses: mudler/localai-github-action@v1.1
|
||||
with:
|
||||
model: 'https://huggingface.co/unsloth/Qwen3.5-2B-GGUF'
|
||||
- name: Process gallery-agent PR commands
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.UPDATE_BOT_TOKEN }}
|
||||
REPO: ${{ github.repository }}
|
||||
SEARCH: 'gallery agent in:title'
|
||||
run: |
|
||||
# Walk gallery-agent PRs and act on maintainer comments:
|
||||
# /gallery-agent blacklist → label `gallery-agent/blacklisted` + close (never repropose)
|
||||
# /gallery-agent recreate → close without label (next run may repropose)
|
||||
# Only comments from OWNER / MEMBER / COLLABORATOR are honored so
|
||||
# random users can't drive the bot.
|
||||
#
|
||||
# We scan both open PRs AND recently-closed PRs that don't already
|
||||
# carry the blacklist label. This covers the common flow where a
|
||||
# maintainer writes /gallery-agent blacklist and immediately clicks
|
||||
# Close — without this, the next scheduled run wouldn't see the
|
||||
# command (PR is already closed) and would repropose the model.
|
||||
gh label create gallery-agent/blacklisted \
|
||||
--repo "$REPO" --color ededed \
|
||||
--description "gallery-agent must not repropose this model" 2>/dev/null || true
|
||||
|
||||
prs_open=$(gh pr list --repo "$REPO" --state open --search "$SEARCH" \
|
||||
--json number --jq '.[].number')
|
||||
# Closed PRs from the last 14 days that don't yet have the blacklist label.
|
||||
# Bounded window keeps the scan cheap while covering late-applied commands.
|
||||
since=$(date -u -d '14 days ago' +%Y-%m-%d)
|
||||
prs_closed=$(gh pr list --repo "$REPO" --state closed \
|
||||
--search "$SEARCH closed:>=$since -label:gallery-agent/blacklisted" \
|
||||
--json number --jq '.[].number')
|
||||
prs=$(printf '%s\n%s\n' "$prs_open" "$prs_closed" | sort -u | sed '/^$/d')
|
||||
for pr in $prs; do
|
||||
state=$(gh pr view "$pr" --repo "$REPO" --json state --jq '.state')
|
||||
cmds=$(gh pr view "$pr" --repo "$REPO" --json comments \
|
||||
--jq '.comments[] | select(.authorAssociation=="OWNER" or .authorAssociation=="MEMBER" or .authorAssociation=="COLLABORATOR") | .body')
|
||||
if echo "$cmds" | grep -qE '(^|[[:space:]])/gallery-agent[[:space:]]+blacklist([[:space:]]|$)'; then
|
||||
echo "PR #$pr: blacklist command found (state=$state)"
|
||||
gh pr edit "$pr" --repo "$REPO" --add-label gallery-agent/blacklisted || true
|
||||
if [ "$state" = "OPEN" ]; then
|
||||
gh pr close "$pr" --repo "$REPO" --comment "Blacklisted via \`/gallery-agent blacklist\`. This model will not be reproposed." || true
|
||||
fi
|
||||
elif [ "$state" = "OPEN" ] && echo "$cmds" | grep -qE '(^|[[:space:]])/gallery-agent[[:space:]]+recreate([[:space:]]|$)'; then
|
||||
echo "PR #$pr: recreate command found"
|
||||
gh pr close "$pr" --repo "$REPO" --comment "Closed via \`/gallery-agent recreate\`. The next scheduled run will propose this model again." || true
|
||||
fi
|
||||
done
|
||||
|
||||
- name: Collect skip URLs for the gallery agent
|
||||
id: open_prs
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
REPO: ${{ github.repository }}
|
||||
SEARCH: 'gallery agent in:title'
|
||||
run: |
|
||||
# Skip set =
|
||||
# URLs from any open gallery-agent PR (avoid duplicate PRs for the same model while one is pending)
|
||||
# + URLs from closed PRs carrying the `gallery-agent/blacklisted` label (hard blacklist)
|
||||
# Plain-closed PRs without the label are ignored — closing a PR is
|
||||
# not by itself a "never propose again" signal; maintainers must
|
||||
# opt in via the /gallery-agent blacklist comment command.
|
||||
urls_open=$(gh pr list --repo "$REPO" --state open --search "$SEARCH" \
|
||||
--json body --jq '[.[].body] | join("\n")' \
|
||||
| grep -oE 'https://huggingface\.co/[^ )]+' || true)
|
||||
urls_blacklist=$(gh pr list --repo "$REPO" --state closed --search "$SEARCH" \
|
||||
--label gallery-agent/blacklisted \
|
||||
--json body --jq '[.[].body] | join("\n")' \
|
||||
| grep -oE 'https://huggingface\.co/[^ )]+' || true)
|
||||
urls=$(printf '%s\n%s\n' "$urls_open" "$urls_blacklist" | sort -u | sed '/^$/d')
|
||||
echo "Skip URLs:"
|
||||
echo "$urls"
|
||||
{
|
||||
echo "urls<<EOF"
|
||||
echo "$urls"
|
||||
echo "EOF"
|
||||
} >> "$GITHUB_OUTPUT"
|
||||
|
||||
- name: Run gallery agent
|
||||
env:
|
||||
#OPENAI_MODEL: ${{ secrets.OPENAI_MODEL }}
|
||||
OPENAI_MODEL: Qwen3.5-2B-GGUF
|
||||
OPENAI_BASE_URL: "http://localhost:8080"
|
||||
OPENAI_KEY: ${{ secrets.OPENAI_KEY }}
|
||||
#OPENAI_BASE_URL: ${{ secrets.OPENAI_BASE_URL }}
|
||||
SEARCH_TERM: ${{ github.event.inputs.search_term || 'GGUF' }}
|
||||
LIMIT: ${{ github.event.inputs.limit || '15' }}
|
||||
QUANTIZATION: ${{ github.event.inputs.quantization || 'Q4_K_M' }}
|
||||
MAX_MODELS: ${{ github.event.inputs.max_models || '1' }}
|
||||
EXTRA_SKIP_URLS: ${{ steps.open_prs.outputs.urls }}
|
||||
run: |
|
||||
export GALLERY_INDEX_PATH=$PWD/gallery/index.yaml
|
||||
go run ./.github/gallery-agent
|
||||
@@ -124,7 +191,21 @@ jobs:
|
||||
|
||||
**Added Models:**
|
||||
${{ steps.read_summary.outputs.added_models || '- No models added' }}
|
||||
|
||||
|
||||
### Bot commands
|
||||
|
||||
Maintainers (owner / member / collaborator) can control this PR
|
||||
by leaving a comment with one of:
|
||||
|
||||
- `/gallery-agent recreate` — close this PR; the next scheduled
|
||||
run will propose this model again (useful if the entry needs
|
||||
to be regenerated with fresh metadata).
|
||||
- `/gallery-agent blacklist` — close this PR and permanently
|
||||
prevent the gallery agent from ever reproposing this model.
|
||||
|
||||
Plain "Close" (without a command) is treated as a no-op: the
|
||||
model may be reproposed by a future run.
|
||||
|
||||
**Workflow Details:**
|
||||
- Triggered by: `${{ github.event_name }}`
|
||||
- Run ID: `${{ github.run_id }}`
|
||||
|
||||
2
.github/workflows/gh-pages.yml
vendored
2
.github/workflows/gh-pages.yml
vendored
@@ -59,7 +59,7 @@ jobs:
|
||||
hugo --minify --baseURL "${{ steps.pages.outputs.base_url }}/"
|
||||
|
||||
- name: Upload artifact
|
||||
uses: actions/upload-pages-artifact@v4
|
||||
uses: actions/upload-pages-artifact@v5
|
||||
with:
|
||||
path: docs/public
|
||||
|
||||
|
||||
4
.github/workflows/release.yaml
vendored
4
.github/workflows/release.yaml
vendored
@@ -39,7 +39,7 @@ jobs:
|
||||
run: |
|
||||
make build-launcher-darwin
|
||||
- name: Upload DMG to Release
|
||||
uses: softprops/action-gh-release@v2
|
||||
uses: softprops/action-gh-release@v3
|
||||
with:
|
||||
files: ./dist/LocalAI.dmg
|
||||
launcher-build-linux:
|
||||
@@ -59,6 +59,6 @@ jobs:
|
||||
sudo apt-get install golang gcc libgl1-mesa-dev xorg-dev libxkbcommon-dev
|
||||
make build-launcher-linux
|
||||
- name: Upload Linux launcher artifacts
|
||||
uses: softprops/action-gh-release@v2
|
||||
uses: softprops/action-gh-release@v3
|
||||
with:
|
||||
files: ./local-ai-launcher-linux.tar.xz
|
||||
|
||||
230
.github/workflows/test-extra.yml
vendored
230
.github/workflows/test-extra.yml
vendored
@@ -31,11 +31,17 @@ jobs:
|
||||
llama-cpp-quantization: ${{ steps.detect.outputs.llama-cpp-quantization }}
|
||||
llama-cpp: ${{ steps.detect.outputs.llama-cpp }}
|
||||
ik-llama-cpp: ${{ steps.detect.outputs.ik-llama-cpp }}
|
||||
turboquant: ${{ steps.detect.outputs.turboquant }}
|
||||
buun-llama-cpp: ${{ steps.detect.outputs['buun-llama-cpp'] }}
|
||||
vllm: ${{ steps.detect.outputs.vllm }}
|
||||
sglang: ${{ steps.detect.outputs.sglang }}
|
||||
acestep-cpp: ${{ steps.detect.outputs.acestep-cpp }}
|
||||
qwen3-tts-cpp: ${{ steps.detect.outputs.qwen3-tts-cpp }}
|
||||
voxtral: ${{ steps.detect.outputs.voxtral }}
|
||||
kokoros: ${{ steps.detect.outputs.kokoros }}
|
||||
insightface: ${{ steps.detect.outputs.insightface }}
|
||||
speaker-recognition: ${{ steps.detect.outputs.speaker-recognition }}
|
||||
sherpa-onnx: ${{ steps.detect.outputs.sherpa-onnx }}
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v6
|
||||
@@ -485,6 +491,89 @@ jobs:
|
||||
- name: Build llama-cpp backend image and run gRPC e2e tests
|
||||
run: |
|
||||
make test-extra-backend-llama-cpp
|
||||
tests-llama-cpp-grpc-transcription:
|
||||
needs: detect-changes
|
||||
if: needs.detect-changes.outputs.llama-cpp == 'true' || needs.detect-changes.outputs.run-all == 'true'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 90
|
||||
steps:
|
||||
- name: Clone
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
submodules: true
|
||||
- name: Setup Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: '1.25.4'
|
||||
- name: Build llama-cpp backend image and run audio transcription gRPC e2e tests
|
||||
run: |
|
||||
make test-extra-backend-llama-cpp-transcription
|
||||
# Realtime e2e with sherpa-onnx driving VAD + STT + TTS against a mocked LLM.
|
||||
# Builds the sherpa-onnx Docker image, extracts the rootfs so the e2e suite
|
||||
# can discover the backend binary + shared libs, downloads the three model
|
||||
# bundles (silero-vad, omnilingual-asr, vits-ljs) and drives the realtime
|
||||
# websocket spec end-to-end.
|
||||
tests-sherpa-onnx-realtime:
|
||||
needs: detect-changes
|
||||
if: needs.detect-changes.outputs.sherpa-onnx == 'true' || needs.detect-changes.outputs.run-all == 'true'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 90
|
||||
steps:
|
||||
- name: Clone
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
submodules: true
|
||||
- name: Setup Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: '1.25.4'
|
||||
- name: Setup Node.js
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
- name: Build sherpa-onnx backend image and run realtime e2e tests
|
||||
run: |
|
||||
make test-extra-e2e-realtime-sherpa
|
||||
# Streaming ASR via the sherpa-onnx online recognizer (zipformer
|
||||
# transducer). Exercises both AudioTranscription (buffered) and
|
||||
# AudioTranscriptionStream (real-time deltas) on the e2e-backends
|
||||
# harness.
|
||||
tests-sherpa-onnx-grpc-transcription:
|
||||
needs: detect-changes
|
||||
if: needs.detect-changes.outputs.sherpa-onnx == 'true' || needs.detect-changes.outputs.run-all == 'true'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 90
|
||||
steps:
|
||||
- name: Clone
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
submodules: true
|
||||
- name: Setup Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: '1.25.4'
|
||||
- name: Build sherpa-onnx backend image and run streaming ASR gRPC e2e tests
|
||||
run: |
|
||||
make test-extra-backend-sherpa-onnx-transcription
|
||||
# VITS TTS via the sherpa-onnx backend. Drives both TTS (file write) and
|
||||
# TTSStream (PCM chunks) on the e2e-backends harness.
|
||||
tests-sherpa-onnx-grpc-tts:
|
||||
needs: detect-changes
|
||||
if: needs.detect-changes.outputs.sherpa-onnx == 'true' || needs.detect-changes.outputs.run-all == 'true'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 90
|
||||
steps:
|
||||
- name: Clone
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
submodules: true
|
||||
- name: Setup Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: '1.25.4'
|
||||
- name: Build sherpa-onnx backend image and run TTS gRPC e2e tests
|
||||
run: |
|
||||
make test-extra-backend-sherpa-onnx-tts
|
||||
tests-ik-llama-cpp-grpc:
|
||||
needs: detect-changes
|
||||
if: needs.detect-changes.outputs.ik-llama-cpp == 'true' || needs.detect-changes.outputs.run-all == 'true'
|
||||
@@ -502,6 +591,53 @@ jobs:
|
||||
- name: Build ik-llama-cpp backend image and run gRPC e2e tests
|
||||
run: |
|
||||
make test-extra-backend-ik-llama-cpp
|
||||
tests-turboquant-grpc:
|
||||
needs: detect-changes
|
||||
if: needs.detect-changes.outputs.turboquant == 'true' || needs.detect-changes.outputs.run-all == 'true'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 90
|
||||
steps:
|
||||
- name: Clone
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
submodules: true
|
||||
- name: Setup Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: '1.25.4'
|
||||
# Exercises the turboquant (llama.cpp fork) backend with KV-cache
|
||||
# quantization enabled. The convenience target sets
|
||||
# BACKEND_TEST_CACHE_TYPE_K / _V=q8_0, which are plumbed into the
|
||||
# ModelOptions.CacheTypeKey/Value gRPC fields. LoadModel-success +
|
||||
# backend stdout/stderr (captured by the Ginkgo suite) prove the
|
||||
# cache-type config path reaches the fork's KV-cache init.
|
||||
- name: Build turboquant backend image and run gRPC e2e tests
|
||||
run: |
|
||||
make test-extra-backend-turboquant
|
||||
tests-buun-llama-cpp-grpc:
|
||||
needs: detect-changes
|
||||
if: needs.detect-changes.outputs['buun-llama-cpp'] == 'true' || needs.detect-changes.outputs.run-all == 'true'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 90
|
||||
steps:
|
||||
- name: Clone
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
submodules: true
|
||||
- name: Setup Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: '1.25.4'
|
||||
# Exercises the buun-llama-cpp (fork-of-a-fork) backend with the
|
||||
# fork-specific TurboQuant/TCQ KV-cache types. BACKEND_TEST_CACHE_TYPE_V
|
||||
# is set to turbo3 so the test round-trips through the fork's KV
|
||||
# allow-list — picking a stock llama.cpp type would only re-test the
|
||||
# shared code path. DFlash speculative decoding is not exercised here
|
||||
# because the one known public target/drafter pair (Qwen3.5-27B) is too
|
||||
# large for CI.
|
||||
- name: Build buun-llama-cpp backend image and run gRPC e2e tests
|
||||
run: |
|
||||
make test-extra-backend-buun-llama-cpp
|
||||
# tests-vllm-grpc is currently disabled in CI.
|
||||
#
|
||||
# The prebuilt vllm CPU wheel is compiled with AVX-512 VNNI/BF16
|
||||
@@ -548,6 +684,48 @@ jobs:
|
||||
# - name: Build vllm (cpu) backend image and run gRPC e2e tests
|
||||
# run: |
|
||||
# make test-extra-backend-vllm
|
||||
# tests-sglang-grpc is currently disabled in CI for the same reason as
|
||||
# tests-vllm-grpc: sglang's CPU kernel (sgl-kernel) uses __m512 AVX-512
|
||||
# intrinsics unconditionally in shm.cpp, so the from-source build
|
||||
# requires `-march=sapphirerapids` (already set in install.sh) and the
|
||||
# resulting binary SIGILLs at import on CPUs without AVX-512 VNNI/BF16.
|
||||
# The ubuntu-latest runner pool does not guarantee that ISA baseline.
|
||||
#
|
||||
# The test itself (tests/e2e-backends + make test-extra-backend-sglang)
|
||||
# is fully working and validated locally on a host with the right
|
||||
# SIMD baseline. Run it manually with:
|
||||
#
|
||||
# make test-extra-backend-sglang
|
||||
#
|
||||
# Re-enable this job once we have a self-hosted runner label with
|
||||
# guaranteed AVX-512 VNNI/BF16 support.
|
||||
#
|
||||
# tests-sglang-grpc:
|
||||
# needs: detect-changes
|
||||
# if: needs.detect-changes.outputs.sglang == 'true' || needs.detect-changes.outputs.run-all == 'true'
|
||||
# runs-on: bigger-runner
|
||||
# timeout-minutes: 90
|
||||
# steps:
|
||||
# - name: Clone
|
||||
# uses: actions/checkout@v6
|
||||
# with:
|
||||
# submodules: true
|
||||
# - name: Dependencies
|
||||
# run: |
|
||||
# sudo apt-get update
|
||||
# sudo apt-get install -y --no-install-recommends \
|
||||
# make build-essential curl unzip ca-certificates git tar
|
||||
# - name: Setup Go
|
||||
# uses: actions/setup-go@v5
|
||||
# with:
|
||||
# go-version: '1.25.4'
|
||||
# - name: Free disk space
|
||||
# run: |
|
||||
# sudo rm -rf /usr/share/dotnet /opt/ghc /usr/local/lib/android /opt/hostedtoolcache/CodeQL || true
|
||||
# df -h
|
||||
# - name: Build sglang (cpu) backend image and run gRPC e2e tests
|
||||
# run: |
|
||||
# make test-extra-backend-sglang
|
||||
tests-acestep-cpp:
|
||||
needs: detect-changes
|
||||
if: needs.detect-changes.outputs.acestep-cpp == 'true' || needs.detect-changes.outputs.run-all == 'true'
|
||||
@@ -667,3 +845,55 @@ jobs:
|
||||
- name: Test kokoros
|
||||
run: |
|
||||
make -C backend/rust/kokoros test
|
||||
tests-insightface-grpc:
|
||||
needs: detect-changes
|
||||
if: needs.detect-changes.outputs.insightface == 'true' || needs.detect-changes.outputs.run-all == 'true'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 90
|
||||
steps:
|
||||
- name: Clone
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
submodules: true
|
||||
- name: Dependencies
|
||||
run: |
|
||||
sudo apt-get update
|
||||
sudo apt-get install -y --no-install-recommends \
|
||||
make build-essential curl unzip ca-certificates git tar
|
||||
- name: Setup Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: '1.26.0'
|
||||
- name: Free disk space
|
||||
run: |
|
||||
sudo rm -rf /usr/share/dotnet /opt/ghc /usr/local/lib/android /opt/hostedtoolcache/CodeQL || true
|
||||
df -h
|
||||
- name: Build insightface backend image and run both model configurations
|
||||
run: |
|
||||
make test-extra-backend-insightface-all
|
||||
tests-speaker-recognition-grpc:
|
||||
needs: detect-changes
|
||||
if: needs.detect-changes.outputs.speaker-recognition == 'true' || needs.detect-changes.outputs.run-all == 'true'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 90
|
||||
steps:
|
||||
- name: Clone
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
submodules: true
|
||||
- name: Dependencies
|
||||
run: |
|
||||
sudo apt-get update
|
||||
sudo apt-get install -y --no-install-recommends \
|
||||
make build-essential curl ca-certificates git tar
|
||||
- name: Setup Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: '1.26.0'
|
||||
- name: Free disk space
|
||||
run: |
|
||||
sudo rm -rf /usr/share/dotnet /opt/ghc /usr/local/lib/android /opt/hostedtoolcache/CodeQL || true
|
||||
df -h
|
||||
- name: Build speaker-recognition backend image and run the ECAPA-TDNN configuration
|
||||
run: |
|
||||
make test-extra-backend-speaker-recognition-all
|
||||
|
||||
2
.github/workflows/test.yml
vendored
2
.github/workflows/test.yml
vendored
@@ -195,7 +195,7 @@ jobs:
|
||||
run: go version
|
||||
- name: Dependencies
|
||||
run: |
|
||||
brew install protobuf grpc make protoc-gen-go protoc-gen-go-grpc libomp llvm opus
|
||||
brew install protobuf grpc make protoc-gen-go protoc-gen-go-grpc libomp llvm opus ffmpeg
|
||||
pip install --user --no-cache-dir grpcio-tools grpcio
|
||||
- name: Setup Node.js
|
||||
uses: actions/setup-node@v6
|
||||
|
||||
18
AGENTS.md
18
AGENTS.md
@@ -1,15 +1,28 @@
|
||||
# LocalAI Agent Instructions
|
||||
|
||||
This file is an index to detailed topic guides in the `.agents/` directory. Read the relevant file(s) for the task at hand — you don't need to load all of them.
|
||||
This file is the entry point for AI coding assistants (Claude Code, Cursor, Copilot, Codex, Aider, etc.) working on LocalAI. It is an index to detailed topic guides in the `.agents/` directory. Read the relevant file(s) for the task at hand — you don't need to load all of them.
|
||||
|
||||
Human contributors: see [CONTRIBUTING.md](CONTRIBUTING.md) for the development workflow.
|
||||
|
||||
## Policy for AI-Assisted Contributions
|
||||
|
||||
LocalAI follows the Linux kernel project's [guidelines for AI coding assistants](https://docs.kernel.org/process/coding-assistants.html). Before submitting AI-assisted code, read [.agents/ai-coding-assistants.md](.agents/ai-coding-assistants.md). Key rules:
|
||||
|
||||
- **No `Signed-off-by` from AI.** Only the human submitter may sign off on the Developer Certificate of Origin.
|
||||
- **No `Co-Authored-By: <AI>` trailers.** The human contributor owns the change.
|
||||
- **Use an `Assisted-by:` trailer** to attribute AI involvement. Format: `Assisted-by: AGENT_NAME:MODEL_VERSION [TOOL1] [TOOL2]`.
|
||||
- **The human submitter is responsible** for reviewing, testing, and understanding every line of generated code.
|
||||
|
||||
## Topics
|
||||
|
||||
| File | When to read |
|
||||
|------|-------------|
|
||||
| [.agents/ai-coding-assistants.md](.agents/ai-coding-assistants.md) | Policy for AI-assisted contributions — licensing, DCO, attribution |
|
||||
| [.agents/building-and-testing.md](.agents/building-and-testing.md) | Building the project, running tests, Docker builds for specific platforms |
|
||||
| [.agents/adding-backends.md](.agents/adding-backends.md) | Adding a new backend (Python, Go, or C++) — full step-by-step checklist |
|
||||
| [.agents/adding-backends.md](.agents/adding-backends.md) | Adding a new backend (Python, Go, or C++) — full step-by-step checklist, including importer integration (the `/import-model` dropdown is server-driven from `GET /backends/known`) |
|
||||
| [.agents/coding-style.md](.agents/coding-style.md) | Code style, editorconfig, logging, documentation conventions |
|
||||
| [.agents/llama-cpp-backend.md](.agents/llama-cpp-backend.md) | Working on the llama.cpp backend — architecture, updating, tool call parsing |
|
||||
| [.agents/vllm-backend.md](.agents/vllm-backend.md) | Working on the vLLM / vLLM-omni backends — native parsers, ChatDelta, CPU build, libnuma packaging, backend hooks |
|
||||
| [.agents/testing-mcp-apps.md](.agents/testing-mcp-apps.md) | Testing MCP Apps (interactive tool UIs) in the React UI |
|
||||
| [.agents/api-endpoints-and-auth.md](.agents/api-endpoints-and-auth.md) | Adding API endpoints, auth middleware, feature permissions, user access control |
|
||||
| [.agents/debugging-backends.md](.agents/debugging-backends.md) | Debugging runtime backend failures, dependency conflicts, rebuilding backends |
|
||||
@@ -21,5 +34,6 @@ This file is an index to detailed topic guides in the `.agents/` directory. Read
|
||||
- **Go style**: Prefer `any` over `interface{}`
|
||||
- **Comments**: Explain *why*, not *what*
|
||||
- **Docs**: Update `docs/content/` when adding features or changing config
|
||||
- **New API endpoints**: LocalAI advertises its capability surface in several independent places — swagger `@Tags`, `/api/instructions` registry, auth `RouteFeatureRegistry`, React UI `capabilities.js`, docs. Read [.agents/api-endpoints-and-auth.md](.agents/api-endpoints-and-auth.md) and follow its checklist — missing any surface means clients, admins, and the UI won't know the endpoint exists.
|
||||
- **Build**: Inspect `Makefile` and `.github/workflows/` — ask the user before running long builds
|
||||
- **UI**: The active UI is the React app in `core/http/react-ui/`. The older Alpine.js/HTML UI in `core/http/static/` is pending deprecation — all new UI work goes in the React UI
|
||||
|
||||
@@ -13,6 +13,7 @@ Thank you for your interest in contributing to LocalAI! We appreciate your time
|
||||
- [Development Workflow](#development-workflow)
|
||||
- [Creating a Pull Request (PR)](#creating-a-pull-request-pr)
|
||||
- [Coding Guidelines](#coding-guidelines)
|
||||
- [AI Coding Assistants](#ai-coding-assistants)
|
||||
- [Testing](#testing)
|
||||
- [Documentation](#documentation)
|
||||
- [Community and Communication](#community-and-communication)
|
||||
@@ -185,7 +186,7 @@ Before jumping into a PR for a massive feature or big change, it is preferred to
|
||||
|
||||
This project uses an [`.editorconfig`](.editorconfig) file to define formatting standards (indentation, line endings, charset, etc.). Please configure your editor to respect it.
|
||||
|
||||
For AI-assisted development, see [`CLAUDE.md`](CLAUDE.md) for agent-specific guidelines including build instructions and backend architecture details.
|
||||
For AI-assisted development, see [`AGENTS.md`](AGENTS.md) (or the equivalent [`CLAUDE.md`](CLAUDE.md) symlink) for agent-specific guidelines including build instructions and backend architecture details. Contributions produced with AI assistance must follow the rules in the [AI Coding Assistants](#ai-coding-assistants) section below.
|
||||
|
||||
### General Principles
|
||||
|
||||
@@ -211,6 +212,26 @@ For AI-assisted development, see [`CLAUDE.md`](CLAUDE.md) for agent-specific gui
|
||||
- Reviewers will check for correctness, test coverage, adherence to these guidelines, and clarity of intent.
|
||||
- Be responsive to review feedback and keep discussions constructive.
|
||||
|
||||
## AI Coding Assistants
|
||||
|
||||
LocalAI follows the **same guidelines as the Linux kernel project** for AI-assisted contributions: <https://docs.kernel.org/process/coding-assistants.html>.
|
||||
|
||||
The full policy for this repository lives in [`.agents/ai-coding-assistants.md`](.agents/ai-coding-assistants.md). Summary:
|
||||
|
||||
- **AI agents MUST NOT add `Signed-off-by` tags.** Only humans can certify the Developer Certificate of Origin.
|
||||
- **AI agents MUST NOT add `Co-Authored-By` trailers** attributing themselves as co-authors.
|
||||
- **Attribute AI involvement with an `Assisted-by` trailer** in the commit message:
|
||||
|
||||
```
|
||||
Assisted-by: AGENT_NAME:MODEL_VERSION [TOOL1] [TOOL2]
|
||||
```
|
||||
|
||||
Example: `Assisted-by: Claude:claude-opus-4-7 golangci-lint`
|
||||
|
||||
Basic development tools (git, go, make, editors) should not be listed.
|
||||
- **The human submitter is responsible** for reviewing, testing, and fully understanding every line of AI-generated code — including verifying that any referenced APIs, flags, or file paths actually exist in the tree.
|
||||
- Contributions must remain compatible with LocalAI's **MIT License**.
|
||||
|
||||
## Testing
|
||||
|
||||
All new features and bug fixes should include test coverage. The project uses [Ginkgo](https://onsi.github.io/ginkgo/) as its test framework.
|
||||
|
||||
375
Makefile
375
Makefile
@@ -1,5 +1,5 @@
|
||||
# Disable parallel execution for backend builds
|
||||
.NOTPARALLEL: backends/diffusers backends/llama-cpp backends/outetts backends/piper backends/stablediffusion-ggml backends/whisper backends/faster-whisper backends/silero-vad backends/local-store backends/huggingface backends/rfdetr backends/kitten-tts backends/kokoro backends/chatterbox backends/llama-cpp-darwin backends/neutts build-darwin-python-backend build-darwin-go-backend backends/mlx backends/diffuser-darwin backends/mlx-vlm backends/mlx-audio backends/mlx-distributed backends/stablediffusion-ggml-darwin backends/vllm backends/vllm-omni backends/moonshine backends/pocket-tts backends/qwen-tts backends/faster-qwen3-tts backends/qwen-asr backends/nemo backends/voxcpm backends/whisperx backends/ace-step backends/acestep-cpp backends/fish-speech backends/voxtral backends/opus backends/trl backends/llama-cpp-quantization backends/kokoros backends/sam3-cpp backends/qwen3-tts-cpp
|
||||
.NOTPARALLEL: backends/diffusers backends/llama-cpp backends/turboquant backends/buun-llama-cpp backends/outetts backends/piper backends/stablediffusion-ggml backends/whisper backends/faster-whisper backends/silero-vad backends/local-store backends/huggingface backends/rfdetr backends/insightface backends/speaker-recognition backends/kitten-tts backends/kokoro backends/chatterbox backends/llama-cpp-darwin backends/neutts build-darwin-python-backend build-darwin-go-backend backends/mlx backends/diffuser-darwin backends/mlx-vlm backends/mlx-audio backends/mlx-distributed backends/stablediffusion-ggml-darwin backends/vllm backends/vllm-omni backends/sglang backends/moonshine backends/pocket-tts backends/qwen-tts backends/faster-qwen3-tts backends/qwen-asr backends/nemo backends/voxcpm backends/whisperx backends/ace-step backends/acestep-cpp backends/fish-speech backends/voxtral backends/opus backends/trl backends/llama-cpp-quantization backends/kokoros backends/sam3-cpp backends/qwen3-tts-cpp backends/tinygrad backends/sherpa-onnx
|
||||
|
||||
GOCMD=go
|
||||
GOTEST=$(GOCMD) test
|
||||
@@ -394,7 +394,13 @@ protoc:
|
||||
.PHONY: protogen-go
|
||||
protogen-go: protoc install-go-tools
|
||||
mkdir -p pkg/grpc/proto
|
||||
./protoc --experimental_allow_proto3_optional -Ibackend/ --go_out=pkg/grpc/proto/ --go_opt=paths=source_relative --go-grpc_out=pkg/grpc/proto/ --go-grpc_opt=paths=source_relative \
|
||||
# install-go-tools writes protoc-gen-go and protoc-gen-go-grpc into
|
||||
# $(shell go env GOPATH)/bin, which isn't on every dev's PATH. protoc
|
||||
# resolves its code-gen plugins via PATH, so without this prefix the
|
||||
# generate step fails with "protoc-gen-go: program not found". Prepend
|
||||
# GOPATH/bin so the freshly-installed plugins win without requiring a
|
||||
# shell-profile change.
|
||||
PATH="$$(go env GOPATH)/bin:$$PATH" ./protoc --experimental_allow_proto3_optional -Ibackend/ --go_out=pkg/grpc/proto/ --go_opt=paths=source_relative --go-grpc_out=pkg/grpc/proto/ --go-grpc_opt=paths=source_relative \
|
||||
backend/backend.proto
|
||||
|
||||
core/config/inference_defaults.json: ## Fetch inference defaults from unsloth (only if missing)
|
||||
@@ -419,6 +425,7 @@ prepare-test-extra: protogen-python
|
||||
$(MAKE) -C backend/python/chatterbox
|
||||
$(MAKE) -C backend/python/vllm
|
||||
$(MAKE) -C backend/python/vllm-omni
|
||||
$(MAKE) -C backend/python/sglang
|
||||
$(MAKE) -C backend/python/vibevoice
|
||||
$(MAKE) -C backend/python/moonshine
|
||||
$(MAKE) -C backend/python/pocket-tts
|
||||
@@ -432,6 +439,9 @@ prepare-test-extra: protogen-python
|
||||
$(MAKE) -C backend/python/whisperx
|
||||
$(MAKE) -C backend/python/ace-step
|
||||
$(MAKE) -C backend/python/trl
|
||||
$(MAKE) -C backend/python/tinygrad
|
||||
$(MAKE) -C backend/python/insightface
|
||||
$(MAKE) -C backend/python/speaker-recognition
|
||||
$(MAKE) -C backend/rust/kokoros kokoros-grpc
|
||||
|
||||
test-extra: prepare-test-extra
|
||||
@@ -454,6 +464,9 @@ test-extra: prepare-test-extra
|
||||
$(MAKE) -C backend/python/whisperx test
|
||||
$(MAKE) -C backend/python/ace-step test
|
||||
$(MAKE) -C backend/python/trl test
|
||||
$(MAKE) -C backend/python/tinygrad test
|
||||
$(MAKE) -C backend/python/insightface test
|
||||
$(MAKE) -C backend/python/speaker-recognition test
|
||||
$(MAKE) -C backend/rust/kokoros test
|
||||
|
||||
##
|
||||
@@ -493,11 +506,24 @@ test-extra-backend: protogen-go
|
||||
BACKEND_TEST_MODEL_URL="$${BACKEND_TEST_MODEL_URL:-$(BACKEND_TEST_MODEL_URL)}" \
|
||||
BACKEND_TEST_MODEL_FILE="$$BACKEND_TEST_MODEL_FILE" \
|
||||
BACKEND_TEST_MODEL_NAME="$$BACKEND_TEST_MODEL_NAME" \
|
||||
BACKEND_TEST_MMPROJ_URL="$$BACKEND_TEST_MMPROJ_URL" \
|
||||
BACKEND_TEST_MMPROJ_FILE="$$BACKEND_TEST_MMPROJ_FILE" \
|
||||
BACKEND_TEST_AUDIO_URL="$$BACKEND_TEST_AUDIO_URL" \
|
||||
BACKEND_TEST_AUDIO_FILE="$$BACKEND_TEST_AUDIO_FILE" \
|
||||
BACKEND_TEST_CAPS="$$BACKEND_TEST_CAPS" \
|
||||
BACKEND_TEST_PROMPT="$$BACKEND_TEST_PROMPT" \
|
||||
BACKEND_TEST_OPTIONS="$$BACKEND_TEST_OPTIONS" \
|
||||
BACKEND_TEST_TOOL_PROMPT="$$BACKEND_TEST_TOOL_PROMPT" \
|
||||
BACKEND_TEST_TOOL_NAME="$$BACKEND_TEST_TOOL_NAME" \
|
||||
BACKEND_TEST_CACHE_TYPE_K="$$BACKEND_TEST_CACHE_TYPE_K" \
|
||||
BACKEND_TEST_CACHE_TYPE_V="$$BACKEND_TEST_CACHE_TYPE_V" \
|
||||
BACKEND_TEST_FACE_IMAGE_1_URL="$$BACKEND_TEST_FACE_IMAGE_1_URL" \
|
||||
BACKEND_TEST_FACE_IMAGE_1_FILE="$$BACKEND_TEST_FACE_IMAGE_1_FILE" \
|
||||
BACKEND_TEST_FACE_IMAGE_2_URL="$$BACKEND_TEST_FACE_IMAGE_2_URL" \
|
||||
BACKEND_TEST_FACE_IMAGE_2_FILE="$$BACKEND_TEST_FACE_IMAGE_2_FILE" \
|
||||
BACKEND_TEST_FACE_IMAGE_3_URL="$$BACKEND_TEST_FACE_IMAGE_3_URL" \
|
||||
BACKEND_TEST_FACE_IMAGE_3_FILE="$$BACKEND_TEST_FACE_IMAGE_3_FILE" \
|
||||
BACKEND_TEST_VERIFY_DISTANCE_CEILING="$$BACKEND_TEST_VERIFY_DISTANCE_CEILING" \
|
||||
go test -v -timeout 30m ./tests/e2e-backends/...
|
||||
|
||||
## Convenience wrappers: build the image, then exercise it.
|
||||
@@ -507,6 +533,44 @@ test-extra-backend-llama-cpp: docker-build-llama-cpp
|
||||
test-extra-backend-ik-llama-cpp: docker-build-ik-llama-cpp
|
||||
BACKEND_IMAGE=local-ai-backend:ik-llama-cpp $(MAKE) test-extra-backend
|
||||
|
||||
## turboquant: exercises the llama.cpp-fork backend with the fork's
|
||||
## *TurboQuant-specific* KV-cache types (turbo3 for both K and V). turbo3
|
||||
## is what makes this backend distinct from stock llama-cpp — picking q8_0
|
||||
## here would only test the standard llama.cpp code path that the upstream
|
||||
## llama-cpp backend already covers. The fork auto-enables flash_attention
|
||||
## when turbo3/turbo4 are active, so we don't need to set it explicitly.
|
||||
test-extra-backend-turboquant: docker-build-turboquant
|
||||
BACKEND_IMAGE=local-ai-backend:turboquant \
|
||||
BACKEND_TEST_CACHE_TYPE_K=q8_0 \
|
||||
BACKEND_TEST_CACHE_TYPE_V=turbo3 \
|
||||
$(MAKE) test-extra-backend
|
||||
|
||||
## buun-llama-cpp: exercises the fork-of-a-fork backend (spiritbuun/buun-llama-cpp)
|
||||
## with the *TurboQuant/TCQ-specific* KV-cache types (turbo3 for V). Same rationale
|
||||
## as turboquant above: picking a standard llama.cpp type would only re-test the
|
||||
## shared code path. buun inherits turboquant's turbo2/turbo3/turbo4 and adds
|
||||
## turbo2_tcq / turbo3_tcq on top. DFlash speculative decoding is not exercised
|
||||
## here because no small DFlash drafter model exists (the known public pair is
|
||||
## Qwen3.5-27B, ~54 GB).
|
||||
test-extra-backend-buun-llama-cpp: docker-build-buun-llama-cpp
|
||||
BACKEND_IMAGE=local-ai-backend:buun-llama-cpp \
|
||||
BACKEND_TEST_CACHE_TYPE_K=q8_0 \
|
||||
BACKEND_TEST_CACHE_TYPE_V=turbo3 \
|
||||
$(MAKE) test-extra-backend
|
||||
|
||||
## Audio transcription wrapper for the llama-cpp backend.
|
||||
## Drives the new AudioTranscription / AudioTranscriptionStream RPCs against
|
||||
## ggml-org/Qwen3-ASR-0.6B-GGUF (a small ASR model that requires its mmproj
|
||||
## audio encoder companion). The audio fixture is a short public-domain
|
||||
## "jfk.wav" clip ggml-org bundles with whisper.cpp's CI assets.
|
||||
test-extra-backend-llama-cpp-transcription: docker-build-llama-cpp
|
||||
BACKEND_IMAGE=local-ai-backend:llama-cpp \
|
||||
BACKEND_TEST_MODEL_URL=https://huggingface.co/ggml-org/Qwen3-ASR-0.6B-GGUF/resolve/main/Qwen3-ASR-0.6B-Q8_0.gguf \
|
||||
BACKEND_TEST_MMPROJ_URL=https://huggingface.co/ggml-org/Qwen3-ASR-0.6B-GGUF/resolve/main/mmproj-Qwen3-ASR-0.6B-Q8_0.gguf \
|
||||
BACKEND_TEST_AUDIO_URL=https://github.com/ggml-org/whisper.cpp/raw/master/samples/jfk.wav \
|
||||
BACKEND_TEST_CAPS=health,load,transcription \
|
||||
$(MAKE) test-extra-backend
|
||||
|
||||
## vllm is resolved from a HuggingFace model id (no file download) and
|
||||
## exercises Predict + streaming + tool-call extraction via the hermes parser.
|
||||
## Requires a host CPU with the SIMD instructions the prebuilt vllm CPU
|
||||
@@ -519,6 +583,287 @@ test-extra-backend-vllm: docker-build-vllm
|
||||
BACKEND_TEST_OPTIONS=tool_parser:hermes \
|
||||
$(MAKE) test-extra-backend
|
||||
|
||||
## tinygrad mirrors the vllm target (same model, same caps, same parser) so
|
||||
## the two backends are directly comparable. The LLM path covers Predict,
|
||||
## streaming and native tool-call extraction. Companion targets below cover
|
||||
## embeddings, Stable Diffusion and Whisper — run them individually or via
|
||||
## the `test-extra-backend-tinygrad-all` aggregate.
|
||||
test-extra-backend-tinygrad: docker-build-tinygrad
|
||||
BACKEND_IMAGE=local-ai-backend:tinygrad \
|
||||
BACKEND_TEST_MODEL_NAME=Qwen/Qwen3-0.6B \
|
||||
BACKEND_TEST_CAPS=health,load,predict,stream,tools \
|
||||
BACKEND_TEST_OPTIONS=tool_parser:hermes \
|
||||
$(MAKE) test-extra-backend
|
||||
|
||||
## tinygrad — embeddings via LLM last-hidden-state pooling. Reuses the same
|
||||
## Qwen3-0.6B as the chat target so we don't need a separate BERT vendor;
|
||||
## the Embedding RPC mean-pools and L2-normalizes the last-layer hidden
|
||||
## state.
|
||||
test-extra-backend-tinygrad-embeddings: docker-build-tinygrad
|
||||
BACKEND_IMAGE=local-ai-backend:tinygrad \
|
||||
BACKEND_TEST_MODEL_NAME=Qwen/Qwen3-0.6B \
|
||||
BACKEND_TEST_CAPS=health,load,embeddings \
|
||||
$(MAKE) test-extra-backend
|
||||
|
||||
## tinygrad — Stable Diffusion 1.5. The original CompVis/runwayml repos have
|
||||
## been gated, so we use the community-maintained mirror at
|
||||
## stable-diffusion-v1-5/stable-diffusion-v1-5 with the EMA-only pruned
|
||||
## checkpoint (~4.3GB). Step count is kept low (4) so a CPU-only run finishes
|
||||
## in a few minutes; bump BACKEND_TEST_IMAGE_STEPS for higher quality.
|
||||
test-extra-backend-tinygrad-sd: docker-build-tinygrad
|
||||
BACKEND_IMAGE=local-ai-backend:tinygrad \
|
||||
BACKEND_TEST_MODEL_NAME=stable-diffusion-v1-5/stable-diffusion-v1-5 \
|
||||
BACKEND_TEST_CAPS=health,load,image \
|
||||
$(MAKE) test-extra-backend
|
||||
|
||||
## tinygrad — Whisper. Loads OpenAI's tiny.en checkpoint (smallest at ~75MB)
|
||||
## from the original azure CDN through tinygrad's `fetch` helper, and
|
||||
## transcribes the canonical jfk.wav fixture from whisper.cpp's CI samples.
|
||||
## Exercises both AudioTranscription and AudioTranscriptionStream.
|
||||
test-extra-backend-tinygrad-whisper: docker-build-tinygrad
|
||||
BACKEND_IMAGE=local-ai-backend:tinygrad \
|
||||
BACKEND_TEST_MODEL_NAME=openai/whisper-tiny.en \
|
||||
BACKEND_TEST_AUDIO_URL=https://github.com/ggml-org/whisper.cpp/raw/master/samples/jfk.wav \
|
||||
BACKEND_TEST_CAPS=health,load,transcription \
|
||||
$(MAKE) test-extra-backend
|
||||
|
||||
test-extra-backend-tinygrad-all: \
|
||||
test-extra-backend-tinygrad \
|
||||
test-extra-backend-tinygrad-embeddings \
|
||||
test-extra-backend-tinygrad-sd \
|
||||
test-extra-backend-tinygrad-whisper
|
||||
|
||||
## insightface — face recognition.
|
||||
##
|
||||
## Face fixtures default to the sample images shipped in the
|
||||
## deepinsight/insightface repository (MIT-licensed). For offline/local
|
||||
## runs override with BACKEND_TEST_FACE_IMAGE_{1,2,3}_FILE pointing at
|
||||
## local paths.
|
||||
FACE_IMAGE_1_URL ?= https://github.com/deepinsight/insightface/raw/master/python-package/insightface/data/images/t1.jpg
|
||||
FACE_IMAGE_2_URL ?= https://github.com/deepinsight/insightface/raw/master/python-package/insightface/data/images/t1.jpg
|
||||
FACE_IMAGE_3_URL ?= https://github.com/deepinsight/insightface/raw/master/python-package/insightface/data/images/mask_white.jpg
|
||||
## Known spoof fixture used by the face_antispoof e2e cap. This is
|
||||
## upstream's own `image_F2.jpg` (Silent-Face repo, via yakhyo mirror)
|
||||
## — verified to classify as is_real=false with score < 0.05 on the
|
||||
## MiniFASNetV2 + MiniFASNetV1SE ensemble.
|
||||
FACE_SPOOF_IMAGE_URL ?= https://github.com/yakhyo/face-anti-spoofing/raw/main/assets/image_F2.jpg
|
||||
|
||||
## Host-side cache for the OpenCV Zoo face ONNX files used by the
|
||||
## opencv e2e target. The backend image no longer bakes model weights —
|
||||
## gallery installs bring them via `files:` — but the e2e suite drives
|
||||
## LoadModel over gRPC directly without going through the gallery. We
|
||||
## pre-download the ONNX files to a stable host path and pass absolute
|
||||
## paths in BACKEND_TEST_OPTIONS; `make` skips the downloads when the
|
||||
## SHA-256 already matches.
|
||||
INSIGHTFACE_OPENCV_DIR := /tmp/localai-insightface-opencv-cache
|
||||
INSIGHTFACE_OPENCV_YUNET_URL := https://github.com/opencv/opencv_zoo/raw/main/models/face_detection_yunet/face_detection_yunet_2023mar.onnx
|
||||
INSIGHTFACE_OPENCV_SFACE_URL := https://github.com/opencv/opencv_zoo/raw/main/models/face_recognition_sface/face_recognition_sface_2021dec.onnx
|
||||
INSIGHTFACE_OPENCV_YUNET_SHA := 8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4
|
||||
INSIGHTFACE_OPENCV_SFACE_SHA := 0ba9fbfa01b5270c96627c4ef784da859931e02f04419c829e83484087c34e79
|
||||
|
||||
## buffalo_sc (insightface) — pack zip + SHA-256 mirrors the gallery
|
||||
## entry so the e2e target matches exactly what `local-ai models install
|
||||
## insightface-buffalo-sc` would have fetched. Smallest insightface pack
|
||||
## (~16MB) — keeps CI fast while still covering the insightface engine
|
||||
## code path end-to-end.
|
||||
INSIGHTFACE_BUFFALO_SC_DIR := /tmp/localai-insightface-buffalo-sc-cache
|
||||
INSIGHTFACE_BUFFALO_SC_URL := https://github.com/deepinsight/insightface/releases/download/v0.7/buffalo_sc.zip
|
||||
INSIGHTFACE_BUFFALO_SC_SHA := 57d31b56b6ffa911c8a73cfc1707c73cab76efe7f13b675a05223bf42de47c72
|
||||
|
||||
## Silent-Face antispoofing (MiniFASNetV2 + MiniFASNetV1SE) — shared
|
||||
## between the buffalo_sc and opencv e2e targets. Both ONNX files are
|
||||
## ~1.7MB, Apache 2.0. URLs + SHAs mirror the gallery entries.
|
||||
INSIGHTFACE_ANTISPOOF_DIR := /tmp/localai-insightface-antispoof-cache
|
||||
INSIGHTFACE_ANTISPOOF_V2_URL := https://github.com/yakhyo/face-anti-spoofing/releases/download/weights/MiniFASNetV2.onnx
|
||||
INSIGHTFACE_ANTISPOOF_V2_SHA := b32929adc2d9c34b9486f8c4c7bc97c1b69bc0ea9befefc380e4faae4e463907
|
||||
INSIGHTFACE_ANTISPOOF_V1SE_URL := https://github.com/yakhyo/face-anti-spoofing/releases/download/weights/MiniFASNetV1SE.onnx
|
||||
INSIGHTFACE_ANTISPOOF_V1SE_SHA := ebab7f90c7833fbccd46d3a555410e78d969db5438e169b6524be444862b3676
|
||||
|
||||
.PHONY: insightface-opencv-models
|
||||
insightface-opencv-models:
|
||||
@mkdir -p $(INSIGHTFACE_OPENCV_DIR)
|
||||
@if [ "$$(sha256sum $(INSIGHTFACE_OPENCV_DIR)/yunet.onnx 2>/dev/null | awk '{print $$1}')" != "$(INSIGHTFACE_OPENCV_YUNET_SHA)" ]; then \
|
||||
echo "Fetching YuNet..."; \
|
||||
curl -fsSL -o $(INSIGHTFACE_OPENCV_DIR)/yunet.onnx $(INSIGHTFACE_OPENCV_YUNET_URL); \
|
||||
echo "$(INSIGHTFACE_OPENCV_YUNET_SHA) $(INSIGHTFACE_OPENCV_DIR)/yunet.onnx" | sha256sum -c; \
|
||||
fi
|
||||
@if [ "$$(sha256sum $(INSIGHTFACE_OPENCV_DIR)/sface.onnx 2>/dev/null | awk '{print $$1}')" != "$(INSIGHTFACE_OPENCV_SFACE_SHA)" ]; then \
|
||||
echo "Fetching SFace..."; \
|
||||
curl -fsSL -o $(INSIGHTFACE_OPENCV_DIR)/sface.onnx $(INSIGHTFACE_OPENCV_SFACE_URL); \
|
||||
echo "$(INSIGHTFACE_OPENCV_SFACE_SHA) $(INSIGHTFACE_OPENCV_DIR)/sface.onnx" | sha256sum -c; \
|
||||
fi
|
||||
|
||||
.PHONY: insightface-antispoof-models
|
||||
insightface-antispoof-models:
|
||||
@mkdir -p $(INSIGHTFACE_ANTISPOOF_DIR)
|
||||
@if [ "$$(sha256sum $(INSIGHTFACE_ANTISPOOF_DIR)/MiniFASNetV2.onnx 2>/dev/null | awk '{print $$1}')" != "$(INSIGHTFACE_ANTISPOOF_V2_SHA)" ]; then \
|
||||
echo "Fetching MiniFASNetV2..."; \
|
||||
curl -fsSL -o $(INSIGHTFACE_ANTISPOOF_DIR)/MiniFASNetV2.onnx $(INSIGHTFACE_ANTISPOOF_V2_URL); \
|
||||
echo "$(INSIGHTFACE_ANTISPOOF_V2_SHA) $(INSIGHTFACE_ANTISPOOF_DIR)/MiniFASNetV2.onnx" | sha256sum -c; \
|
||||
fi
|
||||
@if [ "$$(sha256sum $(INSIGHTFACE_ANTISPOOF_DIR)/MiniFASNetV1SE.onnx 2>/dev/null | awk '{print $$1}')" != "$(INSIGHTFACE_ANTISPOOF_V1SE_SHA)" ]; then \
|
||||
echo "Fetching MiniFASNetV1SE..."; \
|
||||
curl -fsSL -o $(INSIGHTFACE_ANTISPOOF_DIR)/MiniFASNetV1SE.onnx $(INSIGHTFACE_ANTISPOOF_V1SE_URL); \
|
||||
echo "$(INSIGHTFACE_ANTISPOOF_V1SE_SHA) $(INSIGHTFACE_ANTISPOOF_DIR)/MiniFASNetV1SE.onnx" | sha256sum -c; \
|
||||
fi
|
||||
|
||||
.PHONY: insightface-buffalo-sc-models
|
||||
insightface-buffalo-sc-models:
|
||||
@mkdir -p $(INSIGHTFACE_BUFFALO_SC_DIR)
|
||||
@if [ "$$(sha256sum $(INSIGHTFACE_BUFFALO_SC_DIR)/buffalo_sc.zip 2>/dev/null | awk '{print $$1}')" != "$(INSIGHTFACE_BUFFALO_SC_SHA)" ]; then \
|
||||
echo "Fetching buffalo_sc..."; \
|
||||
curl -fsSL -o $(INSIGHTFACE_BUFFALO_SC_DIR)/buffalo_sc.zip $(INSIGHTFACE_BUFFALO_SC_URL); \
|
||||
echo "$(INSIGHTFACE_BUFFALO_SC_SHA) $(INSIGHTFACE_BUFFALO_SC_DIR)/buffalo_sc.zip" | sha256sum -c; \
|
||||
rm -f $(INSIGHTFACE_BUFFALO_SC_DIR)/*.onnx; \
|
||||
fi
|
||||
@if [ ! -f "$(INSIGHTFACE_BUFFALO_SC_DIR)/det_500m.onnx" ]; then \
|
||||
echo "Extracting buffalo_sc..."; \
|
||||
unzip -o -q $(INSIGHTFACE_BUFFALO_SC_DIR)/buffalo_sc.zip -d $(INSIGHTFACE_BUFFALO_SC_DIR); \
|
||||
fi
|
||||
|
||||
## buffalo_sc — smallest insightface pack (SCRFD-500MF detector + MBF
|
||||
## recognizer, ~16MB). Exercises the insightface engine code path
|
||||
## (model_zoo-backed inference) without the ~326MB buffalo_l download.
|
||||
## No age/gender/landmark heads — face_analyze is dropped from caps.
|
||||
## The pack is pre-fetched on the host and passed as `root:<dir>` since
|
||||
## the e2e suite drives LoadModel directly without going through
|
||||
## LocalAI's gallery flow (which is what would normally populate
|
||||
## ModelPath and in turn the engine's `_model_dir` option).
|
||||
test-extra-backend-insightface-buffalo-sc: docker-build-insightface insightface-buffalo-sc-models insightface-antispoof-models
|
||||
BACKEND_IMAGE=local-ai-backend:insightface \
|
||||
BACKEND_TEST_MODEL_NAME=insightface-buffalo-sc \
|
||||
BACKEND_TEST_OPTIONS=engine:insightface,model_pack:buffalo_sc,root:$(INSIGHTFACE_BUFFALO_SC_DIR),antispoof_v2_onnx:$(INSIGHTFACE_ANTISPOOF_DIR)/MiniFASNetV2.onnx,antispoof_v1se_onnx:$(INSIGHTFACE_ANTISPOOF_DIR)/MiniFASNetV1SE.onnx \
|
||||
BACKEND_TEST_CAPS=health,load,face_detect,face_embed,face_verify,face_antispoof \
|
||||
BACKEND_TEST_FACE_IMAGE_1_URL=$(FACE_IMAGE_1_URL) \
|
||||
BACKEND_TEST_FACE_IMAGE_2_URL=$(FACE_IMAGE_2_URL) \
|
||||
BACKEND_TEST_FACE_IMAGE_3_URL=$(FACE_IMAGE_3_URL) \
|
||||
BACKEND_TEST_FACE_SPOOF_IMAGE_URL=$(FACE_SPOOF_IMAGE_URL) \
|
||||
BACKEND_TEST_VERIFY_DISTANCE_CEILING=0.55 \
|
||||
$(MAKE) test-extra-backend
|
||||
|
||||
## OpenCV Zoo YuNet + SFace — Apache 2.0, commercial-safe. face_analyze
|
||||
## cap is dropped (SFace has no demographic head). The ONNX files are
|
||||
## pre-fetched on the host via the insightface-opencv-models target and
|
||||
## passed as absolute paths, since the e2e suite drives LoadModel
|
||||
## directly without going through LocalAI's gallery flow.
|
||||
test-extra-backend-insightface-opencv: docker-build-insightface insightface-opencv-models insightface-antispoof-models
|
||||
BACKEND_IMAGE=local-ai-backend:insightface \
|
||||
BACKEND_TEST_MODEL_NAME=insightface-opencv \
|
||||
BACKEND_TEST_OPTIONS=engine:onnx_direct,detector_onnx:$(INSIGHTFACE_OPENCV_DIR)/yunet.onnx,recognizer_onnx:$(INSIGHTFACE_OPENCV_DIR)/sface.onnx,antispoof_v2_onnx:$(INSIGHTFACE_ANTISPOOF_DIR)/MiniFASNetV2.onnx,antispoof_v1se_onnx:$(INSIGHTFACE_ANTISPOOF_DIR)/MiniFASNetV1SE.onnx \
|
||||
BACKEND_TEST_CAPS=health,load,face_detect,face_embed,face_verify,face_antispoof \
|
||||
BACKEND_TEST_FACE_IMAGE_1_URL=$(FACE_IMAGE_1_URL) \
|
||||
BACKEND_TEST_FACE_IMAGE_2_URL=$(FACE_IMAGE_2_URL) \
|
||||
BACKEND_TEST_FACE_IMAGE_3_URL=$(FACE_IMAGE_3_URL) \
|
||||
BACKEND_TEST_FACE_SPOOF_IMAGE_URL=$(FACE_SPOOF_IMAGE_URL) \
|
||||
BACKEND_TEST_VERIFY_DISTANCE_CEILING=0.55 \
|
||||
$(MAKE) test-extra-backend
|
||||
|
||||
## Aggregate — runs both face-recognition model configurations so CI
|
||||
## catches regressions across engines together.
|
||||
test-extra-backend-insightface-all: \
|
||||
test-extra-backend-insightface-buffalo-sc \
|
||||
test-extra-backend-insightface-opencv
|
||||
|
||||
## speaker-recognition — voice (speaker) biometrics.
|
||||
##
|
||||
## Audio fixtures default to the speechbrain test samples served
|
||||
## straight from their GitHub repo — public, no auth needed, and they
|
||||
## ship as 16kHz mono WAV/FLAC which is exactly what the engine wants.
|
||||
## example{1,2,5} are three different speakers; the suite treats
|
||||
## example1 as the "same-image twin" probe (verify(clip, clip) must
|
||||
## return distance≈0) and the other two as cross-speaker ceilings.
|
||||
## Override with BACKEND_TEST_VOICE_AUDIO_{1,2,3}_FILE for offline runs.
|
||||
VOICE_AUDIO_1_URL ?= https://github.com/speechbrain/speechbrain/raw/develop/tests/samples/single-mic/example1.wav
|
||||
VOICE_AUDIO_2_URL ?= https://github.com/speechbrain/speechbrain/raw/develop/tests/samples/single-mic/example2.flac
|
||||
VOICE_AUDIO_3_URL ?= https://github.com/speechbrain/speechbrain/raw/develop/tests/samples/single-mic/example5.wav
|
||||
|
||||
## ECAPA-TDNN via SpeechBrain — default CI configuration. Auto-downloads
|
||||
## the checkpoint from HuggingFace on first LoadModel (bundled in the
|
||||
## backend image pip install). 192-d embeddings, cosine-distance based.
|
||||
## The e2e suite drives LoadModel directly so we don't rely on LocalAI's
|
||||
## gallery flow here.
|
||||
test-extra-backend-speaker-recognition-ecapa: docker-build-speaker-recognition
|
||||
BACKEND_IMAGE=local-ai-backend:speaker-recognition \
|
||||
BACKEND_TEST_MODEL_NAME=speechbrain/spkrec-ecapa-voxceleb \
|
||||
BACKEND_TEST_OPTIONS=engine:speechbrain,source:speechbrain/spkrec-ecapa-voxceleb \
|
||||
BACKEND_TEST_CAPS=health,load,voice_embed,voice_verify \
|
||||
BACKEND_TEST_VOICE_AUDIO_1_URL=$(VOICE_AUDIO_1_URL) \
|
||||
BACKEND_TEST_VOICE_AUDIO_2_URL=$(VOICE_AUDIO_2_URL) \
|
||||
BACKEND_TEST_VOICE_AUDIO_3_URL=$(VOICE_AUDIO_3_URL) \
|
||||
BACKEND_TEST_VOICE_VERIFY_DISTANCE_CEILING=0.4 \
|
||||
$(MAKE) test-extra-backend
|
||||
|
||||
## Aggregate — today there's only one voice config; the target exists
|
||||
## so the CI workflow matches the insightface-all naming convention and
|
||||
## can grow to include WeSpeaker / 3D-Speaker later.
|
||||
test-extra-backend-speaker-recognition-all: \
|
||||
test-extra-backend-speaker-recognition-ecapa
|
||||
|
||||
## Realtime e2e with sherpa-onnx driving VAD + STT + TTS against a mocked
|
||||
## LLM. Extracts the sherpa-onnx Docker image rootfs, downloads the three
|
||||
## gallery-referenced model bundles (silero-vad, omnilingual-asr, vits-ljs),
|
||||
## writes the corresponding model config YAMLs, and runs the realtime
|
||||
## websocket spec in tests/e2e with REALTIME_* env vars wiring the sherpa
|
||||
## slots into the pipeline. The LLM slot stays on the in-repo mock-backend
|
||||
## registered unconditionally by tests/e2e/e2e_suite_test.go. See
|
||||
## tests/e2e/run-realtime-sherpa.sh for the full orchestration.
|
||||
test-extra-e2e-realtime-sherpa: build-mock-backend docker-build-sherpa-onnx protogen-go react-ui
|
||||
bash tests/e2e/run-realtime-sherpa.sh
|
||||
|
||||
## Streaming ASR via the sherpa-onnx online recognizer. Uses the streaming
|
||||
## zipformer English model (encoder/decoder/joiner int8 + tokens) from the
|
||||
## sherpa-onnx gallery entry. Drives both AudioTranscription and
|
||||
## AudioTranscriptionStream via the e2e-backends gRPC harness; streaming
|
||||
## emits real partial deltas during decode. Each file is renamed on download
|
||||
## to the shape sherpa-onnx's online loader expects (encoder.int8.onnx etc.).
|
||||
test-extra-backend-sherpa-onnx-transcription: docker-build-sherpa-onnx
|
||||
BACKEND_IMAGE=local-ai-backend:sherpa-onnx \
|
||||
BACKEND_TEST_MODEL_URL='https://huggingface.co/csukuangfj/sherpa-onnx-streaming-zipformer-en-2023-06-26/resolve/main/encoder-epoch-99-avg-1-chunk-16-left-128.int8.onnx#encoder.int8.onnx' \
|
||||
BACKEND_TEST_EXTRA_FILES='https://huggingface.co/csukuangfj/sherpa-onnx-streaming-zipformer-en-2023-06-26/resolve/main/decoder-epoch-99-avg-1-chunk-16-left-128.int8.onnx#decoder.int8.onnx|https://huggingface.co/csukuangfj/sherpa-onnx-streaming-zipformer-en-2023-06-26/resolve/main/joiner-epoch-99-avg-1-chunk-16-left-128.int8.onnx#joiner.int8.onnx|https://huggingface.co/csukuangfj/sherpa-onnx-streaming-zipformer-en-2023-06-26/resolve/main/tokens.txt' \
|
||||
BACKEND_TEST_AUDIO_URL=https://github.com/ggml-org/whisper.cpp/raw/master/samples/jfk.wav \
|
||||
BACKEND_TEST_CAPS=health,load,transcription \
|
||||
BACKEND_TEST_OPTIONS=subtype=online \
|
||||
$(MAKE) test-extra-backend
|
||||
|
||||
## VITS TTS via the sherpa-onnx backend. Pulls the individual files from
|
||||
## HuggingFace (the vits-ljs release tarball lives on the k2-fsa github
|
||||
## but is also mirrored as discrete files on HF). Exercises both
|
||||
## TTS (write-to-file) and TTSStream (PCM chunks + WAV header) via the
|
||||
## e2e-backends gRPC harness.
|
||||
test-extra-backend-sherpa-onnx-tts: docker-build-sherpa-onnx
|
||||
BACKEND_IMAGE=local-ai-backend:sherpa-onnx \
|
||||
BACKEND_TEST_MODEL_URL='https://huggingface.co/csukuangfj/vits-ljs/resolve/main/vits-ljs.onnx#vits-ljs.onnx' \
|
||||
BACKEND_TEST_EXTRA_FILES='https://huggingface.co/csukuangfj/vits-ljs/resolve/main/tokens.txt|https://huggingface.co/csukuangfj/vits-ljs/resolve/main/lexicon.txt' \
|
||||
BACKEND_TEST_CAPS=health,load,tts \
|
||||
$(MAKE) test-extra-backend
|
||||
|
||||
## sglang mirrors the vllm setup: HuggingFace model id, same tiny Qwen,
|
||||
## tool-call extraction via sglang's native qwen parser. CPU builds use
|
||||
## sglang's upstream pyproject_cpu.toml recipe (see backend/python/sglang/install.sh).
|
||||
test-extra-backend-sglang: docker-build-sglang
|
||||
BACKEND_IMAGE=local-ai-backend:sglang \
|
||||
BACKEND_TEST_MODEL_NAME=Qwen/Qwen2.5-0.5B-Instruct \
|
||||
BACKEND_TEST_CAPS=health,load,predict,stream,tools \
|
||||
BACKEND_TEST_OPTIONS=tool_parser:qwen \
|
||||
$(MAKE) test-extra-backend
|
||||
|
||||
|
||||
## mlx is Apple-Silicon-first — the MLX backend auto-detects the right tool
|
||||
## parser from the chat template, so no tool_parser: option is needed (it
|
||||
## would be ignored at runtime). Run this on macOS / arm64 with Metal; the
|
||||
## Linux/CPU mlx variant is untested in CI.
|
||||
test-extra-backend-mlx: docker-build-mlx
|
||||
BACKEND_IMAGE=local-ai-backend:mlx \
|
||||
BACKEND_TEST_MODEL_NAME=mlx-community/Qwen2.5-0.5B-Instruct-4bit \
|
||||
BACKEND_TEST_CAPS=health,load,predict,stream,tools \
|
||||
$(MAKE) test-extra-backend
|
||||
|
||||
test-extra-backend-mlx-vlm: docker-build-mlx-vlm
|
||||
BACKEND_IMAGE=local-ai-backend:mlx-vlm \
|
||||
BACKEND_TEST_MODEL_NAME=mlx-community/Qwen2.5-0.5B-Instruct-4bit \
|
||||
BACKEND_TEST_CAPS=health,load,predict,stream,tools \
|
||||
$(MAKE) test-extra-backend
|
||||
|
||||
DOCKER_IMAGE?=local-ai
|
||||
IMAGE_TYPE?=core
|
||||
BASE_IMAGE?=ubuntu:24.04
|
||||
@@ -614,6 +959,14 @@ backend-images:
|
||||
BACKEND_LLAMA_CPP = llama-cpp|llama-cpp|.|false|false
|
||||
# ik-llama-cpp is a fork of llama.cpp with superior CPU performance
|
||||
BACKEND_IK_LLAMA_CPP = ik-llama-cpp|ik-llama-cpp|.|false|false
|
||||
# turboquant is a llama.cpp fork with TurboQuant KV-cache quantization.
|
||||
# Reuses backend/cpp/llama-cpp grpc-server sources via a thin wrapper Makefile.
|
||||
BACKEND_TURBOQUANT = turboquant|turboquant|.|false|false
|
||||
# buun-llama-cpp is a fork-of-a-fork (spiritbuun/buun-llama-cpp forks
|
||||
# TheTom/llama-cpp-turboquant) that adds DFlash block-diffusion speculative
|
||||
# decoding and extra TCQ KV-cache variants on top of TurboQuant. Same thin
|
||||
# wrapper pattern as turboquant — reuses backend/cpp/llama-cpp grpc-server.
|
||||
BACKEND_BUUN_LLAMA_CPP = buun-llama-cpp|buun-llama-cpp|.|false|false
|
||||
|
||||
# Golang backends
|
||||
BACKEND_PIPER = piper|golang|.|false|true
|
||||
@@ -626,6 +979,7 @@ BACKEND_VOXTRAL = voxtral|golang|.|false|true
|
||||
BACKEND_ACESTEP_CPP = acestep-cpp|golang|.|false|true
|
||||
BACKEND_QWEN3_TTS_CPP = qwen3-tts-cpp|golang|.|false|true
|
||||
BACKEND_OPUS = opus|golang|.|false|true
|
||||
BACKEND_SHERPA_ONNX = sherpa-onnx|golang|.|false|true
|
||||
|
||||
# Python backends with root context
|
||||
BACKEND_RERANKERS = rerankers|python|.|false|true
|
||||
@@ -634,11 +988,14 @@ BACKEND_OUTETTS = outetts|python|.|false|true
|
||||
BACKEND_FASTER_WHISPER = faster-whisper|python|.|false|true
|
||||
BACKEND_COQUI = coqui|python|.|false|true
|
||||
BACKEND_RFDETR = rfdetr|python|.|false|true
|
||||
BACKEND_INSIGHTFACE = insightface|python|.|false|true
|
||||
BACKEND_SPEAKER_RECOGNITION = speaker-recognition|python|.|false|true
|
||||
BACKEND_KITTEN_TTS = kitten-tts|python|.|false|true
|
||||
BACKEND_NEUTTS = neutts|python|.|false|true
|
||||
BACKEND_KOKORO = kokoro|python|.|false|true
|
||||
BACKEND_VLLM = vllm|python|.|false|true
|
||||
BACKEND_VLLM_OMNI = vllm-omni|python|.|false|true
|
||||
BACKEND_SGLANG = sglang|python|.|false|true
|
||||
BACKEND_DIFFUSERS = diffusers|python|.|--progress=plain|true
|
||||
BACKEND_CHATTERBOX = chatterbox|python|.|false|true
|
||||
BACKEND_VIBEVOICE = vibevoice|python|.|--progress=plain|true
|
||||
@@ -652,9 +1009,12 @@ BACKEND_NEMO = nemo|python|.|false|true
|
||||
BACKEND_VOXCPM = voxcpm|python|.|false|true
|
||||
BACKEND_WHISPERX = whisperx|python|.|false|true
|
||||
BACKEND_ACE_STEP = ace-step|python|.|false|true
|
||||
BACKEND_MLX = mlx|python|.|false|true
|
||||
BACKEND_MLX_VLM = mlx-vlm|python|.|false|true
|
||||
BACKEND_MLX_DISTRIBUTED = mlx-distributed|python|./|false|true
|
||||
BACKEND_TRL = trl|python|.|false|true
|
||||
BACKEND_LLAMA_CPP_QUANTIZATION = llama-cpp-quantization|python|.|false|true
|
||||
BACKEND_TINYGRAD = tinygrad|python|.|false|true
|
||||
|
||||
# Rust backends
|
||||
BACKEND_KOKOROS = kokoros|rust|.|false|true
|
||||
@@ -686,6 +1046,8 @@ endef
|
||||
# Generate all docker-build targets
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_LLAMA_CPP)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_IK_LLAMA_CPP)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_TURBOQUANT)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_BUUN_LLAMA_CPP)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_PIPER)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_LOCAL_STORE)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_HUGGINGFACE)))
|
||||
@@ -700,11 +1062,14 @@ $(eval $(call generate-docker-build-target,$(BACKEND_OUTETTS)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_FASTER_WHISPER)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_COQUI)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_RFDETR)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_INSIGHTFACE)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_SPEAKER_RECOGNITION)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_KITTEN_TTS)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_NEUTTS)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_KOKORO)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_VLLM)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_VLLM_OMNI)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_SGLANG)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_DIFFUSERS)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_CHATTERBOX)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_VIBEVOICE)))
|
||||
@@ -720,17 +1085,21 @@ $(eval $(call generate-docker-build-target,$(BACKEND_WHISPERX)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_ACE_STEP)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_ACESTEP_CPP)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_QWEN3_TTS_CPP)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_MLX)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_MLX_VLM)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_MLX_DISTRIBUTED)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_TRL)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_LLAMA_CPP_QUANTIZATION)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_TINYGRAD)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_KOKOROS)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_SAM3_CPP)))
|
||||
$(eval $(call generate-docker-build-target,$(BACKEND_SHERPA_ONNX)))
|
||||
|
||||
# Pattern rule for docker-save targets
|
||||
docker-save-%: backend-images
|
||||
docker save local-ai-backend:$* -o backend-images/$*.tar
|
||||
|
||||
docker-build-backends: docker-build-llama-cpp docker-build-ik-llama-cpp docker-build-rerankers docker-build-vllm docker-build-vllm-omni docker-build-transformers docker-build-outetts docker-build-diffusers docker-build-kokoro docker-build-faster-whisper docker-build-coqui docker-build-chatterbox docker-build-vibevoice docker-build-moonshine docker-build-pocket-tts docker-build-qwen-tts docker-build-fish-speech docker-build-faster-qwen3-tts docker-build-qwen-asr docker-build-nemo docker-build-voxcpm docker-build-whisperx docker-build-ace-step docker-build-acestep-cpp docker-build-voxtral docker-build-mlx-distributed docker-build-trl docker-build-llama-cpp-quantization docker-build-kokoros docker-build-sam3-cpp docker-build-qwen3-tts-cpp
|
||||
docker-build-backends: docker-build-llama-cpp docker-build-ik-llama-cpp docker-build-turboquant docker-build-buun-llama-cpp docker-build-rerankers docker-build-vllm docker-build-vllm-omni docker-build-sglang docker-build-transformers docker-build-outetts docker-build-diffusers docker-build-kokoro docker-build-faster-whisper docker-build-coqui docker-build-chatterbox docker-build-vibevoice docker-build-moonshine docker-build-pocket-tts docker-build-qwen-tts docker-build-fish-speech docker-build-faster-qwen3-tts docker-build-qwen-asr docker-build-nemo docker-build-voxcpm docker-build-whisperx docker-build-ace-step docker-build-acestep-cpp docker-build-voxtral docker-build-mlx-distributed docker-build-trl docker-build-llama-cpp-quantization docker-build-tinygrad docker-build-kokoros docker-build-sam3-cpp docker-build-qwen3-tts-cpp docker-build-insightface docker-build-speaker-recognition docker-build-sherpa-onnx
|
||||
|
||||
########################################################
|
||||
### Mock Backend for E2E Tests
|
||||
|
||||
@@ -149,6 +149,7 @@ For more details, see the [Getting Started guide](https://localai.io/basics/gett
|
||||
|
||||
## Latest News
|
||||
|
||||
- **April 2026**: [Voice recognition](https://github.com/mudler/LocalAI/pull/9500), [Face recognition, identification & liveness detection](https://github.com/mudler/LocalAI/pull/9480), [Ollama API compatibility](https://github.com/mudler/LocalAI/pull/9284), [Video generation in stable-diffusion.ggml](https://github.com/mudler/LocalAI/pull/9420), [Backend versioning with auto-upgrade](https://github.com/mudler/LocalAI/pull/9315), [Pin models & load-on-demand toggle](https://github.com/mudler/LocalAI/pull/9309), [Universal model importer](https://github.com/mudler/LocalAI/pull/9466), new backends: [sglang](https://github.com/mudler/LocalAI/pull/9359), [ik-llama-cpp](https://github.com/mudler/LocalAI/pull/9326), [TurboQuant](https://github.com/mudler/LocalAI/pull/9355), [sam.cpp](https://github.com/mudler/LocalAI/pull/9288), [Kokoros](https://github.com/mudler/LocalAI/pull/9212), [qwen3tts.cpp](https://github.com/mudler/LocalAI/pull/9316), [tinygrad multimodal](https://github.com/mudler/LocalAI/pull/9364)
|
||||
- **March 2026**: [Agent management](https://github.com/mudler/LocalAI/pull/8820), [New React UI](https://github.com/mudler/LocalAI/pull/8772), [WebRTC](https://github.com/mudler/LocalAI/pull/8790), [MLX-distributed via P2P and RDMA](https://github.com/mudler/LocalAI/pull/8801), [MCP Apps, MCP Client-side](https://github.com/mudler/LocalAI/pull/8947)
|
||||
- **February 2026**: [Realtime API for audio-to-audio with tool calling](https://github.com/mudler/LocalAI/pull/6245), [ACE-Step 1.5 support](https://github.com/mudler/LocalAI/pull/8396)
|
||||
- **January 2026**: **LocalAI 3.10.0** — Anthropic API support, Open Responses API, video & image generation (LTX-2), unified GPU backends, tool streaming, Moonshine, Pocket-TTS. [Release notes](https://github.com/mudler/LocalAI/releases/tag/v3.10.0)
|
||||
|
||||
290
backend/Dockerfile.buun-llama-cpp
Normal file
290
backend/Dockerfile.buun-llama-cpp
Normal file
@@ -0,0 +1,290 @@
|
||||
ARG BASE_IMAGE=ubuntu:24.04
|
||||
ARG GRPC_BASE_IMAGE=${BASE_IMAGE}
|
||||
|
||||
|
||||
# The grpc target does one thing, it builds and installs GRPC. This is in it's own layer so that it can be effectively cached by CI.
|
||||
# You probably don't need to change anything here, and if you do, make sure that CI is adjusted so that the cache continues to work.
|
||||
FROM ${GRPC_BASE_IMAGE} AS grpc
|
||||
|
||||
# This is a bit of a hack, but it's required in order to be able to effectively cache this layer in CI
|
||||
ARG GRPC_MAKEFLAGS="-j4 -Otarget"
|
||||
ARG GRPC_VERSION=v1.65.0
|
||||
ARG CMAKE_FROM_SOURCE=false
|
||||
# CUDA Toolkit 13.x compatibility: CMake 3.31.9+ fixes toolchain detection/arch table issues
|
||||
ARG CMAKE_VERSION=3.31.10
|
||||
|
||||
ENV MAKEFLAGS=${GRPC_MAKEFLAGS}
|
||||
|
||||
WORKDIR /build
|
||||
|
||||
RUN apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
ca-certificates \
|
||||
build-essential curl libssl-dev \
|
||||
git wget && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install CMake (the version in 22.04 is too old)
|
||||
RUN <<EOT bash
|
||||
if [ "${CMAKE_FROM_SOURCE}" = "true" ]; then
|
||||
curl -L -s https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/cmake-${CMAKE_VERSION}.tar.gz -o cmake.tar.gz && tar xvf cmake.tar.gz && cd cmake-${CMAKE_VERSION} && ./configure && make && make install
|
||||
else
|
||||
apt-get update && \
|
||||
apt-get install -y \
|
||||
cmake && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
fi
|
||||
EOT
|
||||
|
||||
# We install GRPC to a different prefix here so that we can copy in only the build artifacts later
|
||||
# saves several hundred MB on the final docker image size vs copying in the entire GRPC source tree
|
||||
# and running make install in the target container
|
||||
RUN git clone --recurse-submodules --jobs 4 -b ${GRPC_VERSION} --depth 1 --shallow-submodules https://github.com/grpc/grpc && \
|
||||
mkdir -p /build/grpc/cmake/build && \
|
||||
cd /build/grpc/cmake/build && \
|
||||
sed -i "216i\ TESTONLY" "../../third_party/abseil-cpp/absl/container/CMakeLists.txt" && \
|
||||
cmake -DgRPC_INSTALL=ON -DgRPC_BUILD_TESTS=OFF -DCMAKE_INSTALL_PREFIX:PATH=/opt/grpc ../.. && \
|
||||
make && \
|
||||
make install && \
|
||||
rm -rf /build
|
||||
|
||||
FROM ${BASE_IMAGE} AS builder
|
||||
ARG CMAKE_FROM_SOURCE=false
|
||||
ARG CMAKE_VERSION=3.31.10
|
||||
# We can target specific CUDA ARCHITECTURES like --build-arg CUDA_DOCKER_ARCH='75;86;89;120'
|
||||
ARG CUDA_DOCKER_ARCH
|
||||
ENV CUDA_DOCKER_ARCH=${CUDA_DOCKER_ARCH}
|
||||
ARG CMAKE_ARGS
|
||||
ENV CMAKE_ARGS=${CMAKE_ARGS}
|
||||
ARG BACKEND=rerankers
|
||||
ARG BUILD_TYPE
|
||||
ENV BUILD_TYPE=${BUILD_TYPE}
|
||||
ARG CUDA_MAJOR_VERSION
|
||||
ARG CUDA_MINOR_VERSION
|
||||
ARG SKIP_DRIVERS=false
|
||||
ENV CUDA_MAJOR_VERSION=${CUDA_MAJOR_VERSION}
|
||||
ENV CUDA_MINOR_VERSION=${CUDA_MINOR_VERSION}
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
ARG TARGETARCH
|
||||
ARG TARGETVARIANT
|
||||
ARG GO_VERSION=1.25.4
|
||||
ARG UBUNTU_VERSION=2404
|
||||
|
||||
RUN apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
build-essential \
|
||||
ccache git \
|
||||
ca-certificates \
|
||||
make \
|
||||
pkg-config libcurl4-openssl-dev \
|
||||
curl unzip \
|
||||
libssl-dev wget && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Cuda
|
||||
ENV PATH=/usr/local/cuda/bin:${PATH}
|
||||
|
||||
# HipBLAS requirements
|
||||
ENV PATH=/opt/rocm/bin:${PATH}
|
||||
|
||||
|
||||
# Vulkan requirements
|
||||
RUN <<EOT bash
|
||||
if [ "${BUILD_TYPE}" = "vulkan" ] && [ "${SKIP_DRIVERS}" = "false" ]; then
|
||||
apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
software-properties-common pciutils wget gpg-agent && \
|
||||
apt-get install -y libglm-dev cmake libxcb-dri3-0 libxcb-present0 libpciaccess0 \
|
||||
libpng-dev libxcb-keysyms1-dev libxcb-dri3-dev libx11-dev g++ gcc \
|
||||
libwayland-dev libxrandr-dev libxcb-randr0-dev libxcb-ewmh-dev \
|
||||
git python-is-python3 bison libx11-xcb-dev liblz4-dev libzstd-dev \
|
||||
ocaml-core ninja-build pkg-config libxml2-dev wayland-protocols python3-jsonschema \
|
||||
clang-format qtbase5-dev qt6-base-dev libxcb-glx0-dev sudo xz-utils
|
||||
if [ "amd64" = "$TARGETARCH" ]; then
|
||||
wget "https://sdk.lunarg.com/sdk/download/1.4.335.0/linux/vulkansdk-linux-x86_64-1.4.335.0.tar.xz" && \
|
||||
tar -xf vulkansdk-linux-x86_64-1.4.335.0.tar.xz && \
|
||||
rm vulkansdk-linux-x86_64-1.4.335.0.tar.xz && \
|
||||
mkdir -p /opt/vulkan-sdk && \
|
||||
mv 1.4.335.0 /opt/vulkan-sdk/ && \
|
||||
cd /opt/vulkan-sdk/1.4.335.0 && \
|
||||
./vulkansdk --no-deps --maxjobs \
|
||||
vulkan-loader \
|
||||
vulkan-validationlayers \
|
||||
vulkan-extensionlayer \
|
||||
vulkan-tools \
|
||||
shaderc && \
|
||||
cp -rfv /opt/vulkan-sdk/1.4.335.0/x86_64/bin/* /usr/bin/ && \
|
||||
cp -rfv /opt/vulkan-sdk/1.4.335.0/x86_64/lib/* /usr/lib/x86_64-linux-gnu/ && \
|
||||
cp -rfv /opt/vulkan-sdk/1.4.335.0/x86_64/include/* /usr/include/ && \
|
||||
cp -rfv /opt/vulkan-sdk/1.4.335.0/x86_64/share/* /usr/share/ && \
|
||||
rm -rf /opt/vulkan-sdk
|
||||
fi
|
||||
if [ "arm64" = "$TARGETARCH" ]; then
|
||||
mkdir vulkan && cd vulkan && \
|
||||
curl -L -o vulkan-sdk.tar.xz https://github.com/mudler/vulkan-sdk-arm/releases/download/1.4.335.0/vulkansdk-ubuntu-24.04-arm-1.4.335.0.tar.xz && \
|
||||
tar -xvf vulkan-sdk.tar.xz && \
|
||||
rm vulkan-sdk.tar.xz && \
|
||||
cd 1.4.335.0 && \
|
||||
cp -rfv aarch64/bin/* /usr/bin/ && \
|
||||
cp -rfv aarch64/lib/* /usr/lib/aarch64-linux-gnu/ && \
|
||||
cp -rfv aarch64/include/* /usr/include/ && \
|
||||
cp -rfv aarch64/share/* /usr/share/ && \
|
||||
cd ../.. && \
|
||||
rm -rf vulkan
|
||||
fi
|
||||
ldconfig && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
fi
|
||||
EOT
|
||||
|
||||
# CuBLAS requirements
|
||||
RUN <<EOT bash
|
||||
if ( [ "${BUILD_TYPE}" = "cublas" ] || [ "${BUILD_TYPE}" = "l4t" ] ) && [ "${SKIP_DRIVERS}" = "false" ]; then
|
||||
apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
software-properties-common pciutils
|
||||
if [ "amd64" = "$TARGETARCH" ]; then
|
||||
curl -O https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION}/x86_64/cuda-keyring_1.1-1_all.deb
|
||||
fi
|
||||
if [ "arm64" = "$TARGETARCH" ]; then
|
||||
if [ "${CUDA_MAJOR_VERSION}" = "13" ]; then
|
||||
curl -O https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION}/sbsa/cuda-keyring_1.1-1_all.deb
|
||||
else
|
||||
curl -O https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION}/arm64/cuda-keyring_1.1-1_all.deb
|
||||
fi
|
||||
fi
|
||||
dpkg -i cuda-keyring_1.1-1_all.deb && \
|
||||
rm -f cuda-keyring_1.1-1_all.deb && \
|
||||
apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
cuda-nvcc-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \
|
||||
libcufft-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \
|
||||
libcurand-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \
|
||||
libcublas-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \
|
||||
libcusparse-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \
|
||||
libcusolver-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION}
|
||||
if [ "${CUDA_MAJOR_VERSION}" = "13" ] && [ "arm64" = "$TARGETARCH" ]; then
|
||||
apt-get install -y --no-install-recommends \
|
||||
libcufile-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} libcudnn9-cuda-${CUDA_MAJOR_VERSION} cuda-cupti-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} libnvjitlink-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION}
|
||||
fi
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
fi
|
||||
EOT
|
||||
|
||||
|
||||
# https://github.com/NVIDIA/Isaac-GR00T/issues/343
|
||||
RUN <<EOT bash
|
||||
if [ "${BUILD_TYPE}" = "cublas" ] && [ "${TARGETARCH}" = "arm64" ]; then
|
||||
wget https://developer.download.nvidia.com/compute/cudss/0.6.0/local_installers/cudss-local-tegra-repo-ubuntu${UBUNTU_VERSION}-0.6.0_0.6.0-1_arm64.deb && \
|
||||
dpkg -i cudss-local-tegra-repo-ubuntu${UBUNTU_VERSION}-0.6.0_0.6.0-1_arm64.deb && \
|
||||
cp /var/cudss-local-tegra-repo-ubuntu${UBUNTU_VERSION}-0.6.0/cudss-*-keyring.gpg /usr/share/keyrings/ && \
|
||||
apt-get update && apt-get -y install cudss cudss-cuda-${CUDA_MAJOR_VERSION} && \
|
||||
wget https://developer.download.nvidia.com/compute/nvpl/25.5/local_installers/nvpl-local-repo-ubuntu${UBUNTU_VERSION}-25.5_1.0-1_arm64.deb && \
|
||||
dpkg -i nvpl-local-repo-ubuntu${UBUNTU_VERSION}-25.5_1.0-1_arm64.deb && \
|
||||
cp /var/nvpl-local-repo-ubuntu${UBUNTU_VERSION}-25.5/nvpl-*-keyring.gpg /usr/share/keyrings/ && \
|
||||
apt-get update && apt-get install -y nvpl
|
||||
fi
|
||||
EOT
|
||||
|
||||
# If we are building with clblas support, we need the libraries for the builds
|
||||
RUN if [ "${BUILD_TYPE}" = "clblas" ] && [ "${SKIP_DRIVERS}" = "false" ]; then \
|
||||
apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
libclblast-dev && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/* \
|
||||
; fi
|
||||
|
||||
RUN if [ "${BUILD_TYPE}" = "hipblas" ] && [ "${SKIP_DRIVERS}" = "false" ]; then \
|
||||
apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
hipblas-dev \
|
||||
rocblas-dev && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/* && \
|
||||
# I have no idea why, but the ROCM lib packages don't trigger ldconfig after they install, which results in local-ai and others not being able
|
||||
# to locate the libraries. We run ldconfig ourselves to work around this packaging deficiency
|
||||
ldconfig && \
|
||||
# Log which GPU architectures have rocBLAS kernel support
|
||||
echo "rocBLAS library data architectures:" && \
|
||||
(ls /opt/rocm*/lib/rocblas/library/Kernels* 2>/dev/null || ls /opt/rocm*/lib64/rocblas/library/Kernels* 2>/dev/null) | grep -oP 'gfx[0-9a-z+-]+' | sort -u || \
|
||||
echo "WARNING: No rocBLAS kernel data found" \
|
||||
; fi
|
||||
|
||||
RUN echo "TARGETARCH: $TARGETARCH"
|
||||
|
||||
# We need protoc installed, and the version in 22.04 is too old. We will create one as part installing the GRPC build below
|
||||
# but that will also being in a newer version of absl which stablediffusion cannot compile with. This version of protoc is only
|
||||
# here so that we can generate the grpc code for the stablediffusion build
|
||||
RUN <<EOT bash
|
||||
if [ "amd64" = "$TARGETARCH" ]; then
|
||||
curl -L -s https://github.com/protocolbuffers/protobuf/releases/download/v27.1/protoc-27.1-linux-x86_64.zip -o protoc.zip && \
|
||||
unzip -j -d /usr/local/bin protoc.zip bin/protoc && \
|
||||
rm protoc.zip
|
||||
fi
|
||||
if [ "arm64" = "$TARGETARCH" ]; then
|
||||
curl -L -s https://github.com/protocolbuffers/protobuf/releases/download/v27.1/protoc-27.1-linux-aarch_64.zip -o protoc.zip && \
|
||||
unzip -j -d /usr/local/bin protoc.zip bin/protoc && \
|
||||
rm protoc.zip
|
||||
fi
|
||||
EOT
|
||||
|
||||
# Install CMake (the version in 22.04 is too old)
|
||||
RUN <<EOT bash
|
||||
if [ "${CMAKE_FROM_SOURCE}" = "true" ]; then
|
||||
curl -L -s https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/cmake-${CMAKE_VERSION}.tar.gz -o cmake.tar.gz && tar xvf cmake.tar.gz && cd cmake-${CMAKE_VERSION} && ./configure && make && make install
|
||||
else
|
||||
apt-get update && \
|
||||
apt-get install -y \
|
||||
cmake && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
fi
|
||||
EOT
|
||||
|
||||
COPY --from=grpc /opt/grpc /usr/local
|
||||
|
||||
|
||||
COPY . /LocalAI
|
||||
|
||||
RUN <<'EOT' bash
|
||||
set -euxo pipefail
|
||||
|
||||
if [[ -n "${CUDA_DOCKER_ARCH:-}" ]]; then
|
||||
CUDA_ARCH_ESC="${CUDA_DOCKER_ARCH//;/\\;}"
|
||||
export CMAKE_ARGS="${CMAKE_ARGS:-} -DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCH_ESC}"
|
||||
echo "CMAKE_ARGS(env) = ${CMAKE_ARGS}"
|
||||
rm -rf /LocalAI/backend/cpp/buun-llama-cpp-*-build
|
||||
fi
|
||||
|
||||
cd /LocalAI/backend/cpp/buun-llama-cpp
|
||||
|
||||
if [ "${TARGETARCH}" = "arm64" ] || [ "${BUILD_TYPE}" = "hipblas" ]; then
|
||||
make buun-llama-cpp-fallback
|
||||
make buun-llama-cpp-grpc
|
||||
make buun-llama-cpp-rpc-server
|
||||
else
|
||||
make buun-llama-cpp-avx
|
||||
make buun-llama-cpp-avx2
|
||||
make buun-llama-cpp-avx512
|
||||
make buun-llama-cpp-fallback
|
||||
make buun-llama-cpp-grpc
|
||||
make buun-llama-cpp-rpc-server
|
||||
fi
|
||||
EOT
|
||||
|
||||
|
||||
# Copy libraries using a script to handle architecture differences
|
||||
RUN make -BC /LocalAI/backend/cpp/buun-llama-cpp package
|
||||
|
||||
|
||||
FROM scratch
|
||||
|
||||
|
||||
# Copy all available binaries (the build process only creates the appropriate ones for the target architecture)
|
||||
COPY --from=builder /LocalAI/backend/cpp/buun-llama-cpp/package/. ./
|
||||
@@ -58,6 +58,8 @@ ARG CUDA_DOCKER_ARCH
|
||||
ENV CUDA_DOCKER_ARCH=${CUDA_DOCKER_ARCH}
|
||||
ARG CMAKE_ARGS
|
||||
ENV CMAKE_ARGS=${CMAKE_ARGS}
|
||||
ARG AMDGPU_TARGETS
|
||||
ENV AMDGPU_TARGETS=${AMDGPU_TARGETS}
|
||||
ARG BACKEND=rerankers
|
||||
ARG BUILD_TYPE
|
||||
ENV BUILD_TYPE=${BUILD_TYPE}
|
||||
|
||||
290
backend/Dockerfile.turboquant
Normal file
290
backend/Dockerfile.turboquant
Normal file
@@ -0,0 +1,290 @@
|
||||
ARG BASE_IMAGE=ubuntu:24.04
|
||||
ARG GRPC_BASE_IMAGE=${BASE_IMAGE}
|
||||
|
||||
|
||||
# The grpc target does one thing, it builds and installs GRPC. This is in it's own layer so that it can be effectively cached by CI.
|
||||
# You probably don't need to change anything here, and if you do, make sure that CI is adjusted so that the cache continues to work.
|
||||
FROM ${GRPC_BASE_IMAGE} AS grpc
|
||||
|
||||
# This is a bit of a hack, but it's required in order to be able to effectively cache this layer in CI
|
||||
ARG GRPC_MAKEFLAGS="-j4 -Otarget"
|
||||
ARG GRPC_VERSION=v1.65.0
|
||||
ARG CMAKE_FROM_SOURCE=false
|
||||
# CUDA Toolkit 13.x compatibility: CMake 3.31.9+ fixes toolchain detection/arch table issues
|
||||
ARG CMAKE_VERSION=3.31.10
|
||||
|
||||
ENV MAKEFLAGS=${GRPC_MAKEFLAGS}
|
||||
|
||||
WORKDIR /build
|
||||
|
||||
RUN apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
ca-certificates \
|
||||
build-essential curl libssl-dev \
|
||||
git wget && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install CMake (the version in 22.04 is too old)
|
||||
RUN <<EOT bash
|
||||
if [ "${CMAKE_FROM_SOURCE}" = "true" ]; then
|
||||
curl -L -s https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/cmake-${CMAKE_VERSION}.tar.gz -o cmake.tar.gz && tar xvf cmake.tar.gz && cd cmake-${CMAKE_VERSION} && ./configure && make && make install
|
||||
else
|
||||
apt-get update && \
|
||||
apt-get install -y \
|
||||
cmake && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
fi
|
||||
EOT
|
||||
|
||||
# We install GRPC to a different prefix here so that we can copy in only the build artifacts later
|
||||
# saves several hundred MB on the final docker image size vs copying in the entire GRPC source tree
|
||||
# and running make install in the target container
|
||||
RUN git clone --recurse-submodules --jobs 4 -b ${GRPC_VERSION} --depth 1 --shallow-submodules https://github.com/grpc/grpc && \
|
||||
mkdir -p /build/grpc/cmake/build && \
|
||||
cd /build/grpc/cmake/build && \
|
||||
sed -i "216i\ TESTONLY" "../../third_party/abseil-cpp/absl/container/CMakeLists.txt" && \
|
||||
cmake -DgRPC_INSTALL=ON -DgRPC_BUILD_TESTS=OFF -DCMAKE_INSTALL_PREFIX:PATH=/opt/grpc ../.. && \
|
||||
make && \
|
||||
make install && \
|
||||
rm -rf /build
|
||||
|
||||
FROM ${BASE_IMAGE} AS builder
|
||||
ARG CMAKE_FROM_SOURCE=false
|
||||
ARG CMAKE_VERSION=3.31.10
|
||||
# We can target specific CUDA ARCHITECTURES like --build-arg CUDA_DOCKER_ARCH='75;86;89;120'
|
||||
ARG CUDA_DOCKER_ARCH
|
||||
ENV CUDA_DOCKER_ARCH=${CUDA_DOCKER_ARCH}
|
||||
ARG CMAKE_ARGS
|
||||
ENV CMAKE_ARGS=${CMAKE_ARGS}
|
||||
ARG BACKEND=rerankers
|
||||
ARG BUILD_TYPE
|
||||
ENV BUILD_TYPE=${BUILD_TYPE}
|
||||
ARG CUDA_MAJOR_VERSION
|
||||
ARG CUDA_MINOR_VERSION
|
||||
ARG SKIP_DRIVERS=false
|
||||
ENV CUDA_MAJOR_VERSION=${CUDA_MAJOR_VERSION}
|
||||
ENV CUDA_MINOR_VERSION=${CUDA_MINOR_VERSION}
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
ARG TARGETARCH
|
||||
ARG TARGETVARIANT
|
||||
ARG GO_VERSION=1.25.4
|
||||
ARG UBUNTU_VERSION=2404
|
||||
|
||||
RUN apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
build-essential \
|
||||
ccache git \
|
||||
ca-certificates \
|
||||
make \
|
||||
pkg-config libcurl4-openssl-dev \
|
||||
curl unzip \
|
||||
libssl-dev wget && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Cuda
|
||||
ENV PATH=/usr/local/cuda/bin:${PATH}
|
||||
|
||||
# HipBLAS requirements
|
||||
ENV PATH=/opt/rocm/bin:${PATH}
|
||||
|
||||
|
||||
# Vulkan requirements
|
||||
RUN <<EOT bash
|
||||
if [ "${BUILD_TYPE}" = "vulkan" ] && [ "${SKIP_DRIVERS}" = "false" ]; then
|
||||
apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
software-properties-common pciutils wget gpg-agent && \
|
||||
apt-get install -y libglm-dev cmake libxcb-dri3-0 libxcb-present0 libpciaccess0 \
|
||||
libpng-dev libxcb-keysyms1-dev libxcb-dri3-dev libx11-dev g++ gcc \
|
||||
libwayland-dev libxrandr-dev libxcb-randr0-dev libxcb-ewmh-dev \
|
||||
git python-is-python3 bison libx11-xcb-dev liblz4-dev libzstd-dev \
|
||||
ocaml-core ninja-build pkg-config libxml2-dev wayland-protocols python3-jsonschema \
|
||||
clang-format qtbase5-dev qt6-base-dev libxcb-glx0-dev sudo xz-utils
|
||||
if [ "amd64" = "$TARGETARCH" ]; then
|
||||
wget "https://sdk.lunarg.com/sdk/download/1.4.335.0/linux/vulkansdk-linux-x86_64-1.4.335.0.tar.xz" && \
|
||||
tar -xf vulkansdk-linux-x86_64-1.4.335.0.tar.xz && \
|
||||
rm vulkansdk-linux-x86_64-1.4.335.0.tar.xz && \
|
||||
mkdir -p /opt/vulkan-sdk && \
|
||||
mv 1.4.335.0 /opt/vulkan-sdk/ && \
|
||||
cd /opt/vulkan-sdk/1.4.335.0 && \
|
||||
./vulkansdk --no-deps --maxjobs \
|
||||
vulkan-loader \
|
||||
vulkan-validationlayers \
|
||||
vulkan-extensionlayer \
|
||||
vulkan-tools \
|
||||
shaderc && \
|
||||
cp -rfv /opt/vulkan-sdk/1.4.335.0/x86_64/bin/* /usr/bin/ && \
|
||||
cp -rfv /opt/vulkan-sdk/1.4.335.0/x86_64/lib/* /usr/lib/x86_64-linux-gnu/ && \
|
||||
cp -rfv /opt/vulkan-sdk/1.4.335.0/x86_64/include/* /usr/include/ && \
|
||||
cp -rfv /opt/vulkan-sdk/1.4.335.0/x86_64/share/* /usr/share/ && \
|
||||
rm -rf /opt/vulkan-sdk
|
||||
fi
|
||||
if [ "arm64" = "$TARGETARCH" ]; then
|
||||
mkdir vulkan && cd vulkan && \
|
||||
curl -L -o vulkan-sdk.tar.xz https://github.com/mudler/vulkan-sdk-arm/releases/download/1.4.335.0/vulkansdk-ubuntu-24.04-arm-1.4.335.0.tar.xz && \
|
||||
tar -xvf vulkan-sdk.tar.xz && \
|
||||
rm vulkan-sdk.tar.xz && \
|
||||
cd 1.4.335.0 && \
|
||||
cp -rfv aarch64/bin/* /usr/bin/ && \
|
||||
cp -rfv aarch64/lib/* /usr/lib/aarch64-linux-gnu/ && \
|
||||
cp -rfv aarch64/include/* /usr/include/ && \
|
||||
cp -rfv aarch64/share/* /usr/share/ && \
|
||||
cd ../.. && \
|
||||
rm -rf vulkan
|
||||
fi
|
||||
ldconfig && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
fi
|
||||
EOT
|
||||
|
||||
# CuBLAS requirements
|
||||
RUN <<EOT bash
|
||||
if ( [ "${BUILD_TYPE}" = "cublas" ] || [ "${BUILD_TYPE}" = "l4t" ] ) && [ "${SKIP_DRIVERS}" = "false" ]; then
|
||||
apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
software-properties-common pciutils
|
||||
if [ "amd64" = "$TARGETARCH" ]; then
|
||||
curl -O https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION}/x86_64/cuda-keyring_1.1-1_all.deb
|
||||
fi
|
||||
if [ "arm64" = "$TARGETARCH" ]; then
|
||||
if [ "${CUDA_MAJOR_VERSION}" = "13" ]; then
|
||||
curl -O https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION}/sbsa/cuda-keyring_1.1-1_all.deb
|
||||
else
|
||||
curl -O https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION}/arm64/cuda-keyring_1.1-1_all.deb
|
||||
fi
|
||||
fi
|
||||
dpkg -i cuda-keyring_1.1-1_all.deb && \
|
||||
rm -f cuda-keyring_1.1-1_all.deb && \
|
||||
apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
cuda-nvcc-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \
|
||||
libcufft-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \
|
||||
libcurand-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \
|
||||
libcublas-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \
|
||||
libcusparse-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \
|
||||
libcusolver-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION}
|
||||
if [ "${CUDA_MAJOR_VERSION}" = "13" ] && [ "arm64" = "$TARGETARCH" ]; then
|
||||
apt-get install -y --no-install-recommends \
|
||||
libcufile-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} libcudnn9-cuda-${CUDA_MAJOR_VERSION} cuda-cupti-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} libnvjitlink-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION}
|
||||
fi
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
fi
|
||||
EOT
|
||||
|
||||
|
||||
# https://github.com/NVIDIA/Isaac-GR00T/issues/343
|
||||
RUN <<EOT bash
|
||||
if [ "${BUILD_TYPE}" = "cublas" ] && [ "${TARGETARCH}" = "arm64" ]; then
|
||||
wget https://developer.download.nvidia.com/compute/cudss/0.6.0/local_installers/cudss-local-tegra-repo-ubuntu${UBUNTU_VERSION}-0.6.0_0.6.0-1_arm64.deb && \
|
||||
dpkg -i cudss-local-tegra-repo-ubuntu${UBUNTU_VERSION}-0.6.0_0.6.0-1_arm64.deb && \
|
||||
cp /var/cudss-local-tegra-repo-ubuntu${UBUNTU_VERSION}-0.6.0/cudss-*-keyring.gpg /usr/share/keyrings/ && \
|
||||
apt-get update && apt-get -y install cudss cudss-cuda-${CUDA_MAJOR_VERSION} && \
|
||||
wget https://developer.download.nvidia.com/compute/nvpl/25.5/local_installers/nvpl-local-repo-ubuntu${UBUNTU_VERSION}-25.5_1.0-1_arm64.deb && \
|
||||
dpkg -i nvpl-local-repo-ubuntu${UBUNTU_VERSION}-25.5_1.0-1_arm64.deb && \
|
||||
cp /var/nvpl-local-repo-ubuntu${UBUNTU_VERSION}-25.5/nvpl-*-keyring.gpg /usr/share/keyrings/ && \
|
||||
apt-get update && apt-get install -y nvpl
|
||||
fi
|
||||
EOT
|
||||
|
||||
# If we are building with clblas support, we need the libraries for the builds
|
||||
RUN if [ "${BUILD_TYPE}" = "clblas" ] && [ "${SKIP_DRIVERS}" = "false" ]; then \
|
||||
apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
libclblast-dev && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/* \
|
||||
; fi
|
||||
|
||||
RUN if [ "${BUILD_TYPE}" = "hipblas" ] && [ "${SKIP_DRIVERS}" = "false" ]; then \
|
||||
apt-get update && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
hipblas-dev \
|
||||
rocblas-dev && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/* && \
|
||||
# I have no idea why, but the ROCM lib packages don't trigger ldconfig after they install, which results in local-ai and others not being able
|
||||
# to locate the libraries. We run ldconfig ourselves to work around this packaging deficiency
|
||||
ldconfig && \
|
||||
# Log which GPU architectures have rocBLAS kernel support
|
||||
echo "rocBLAS library data architectures:" && \
|
||||
(ls /opt/rocm*/lib/rocblas/library/Kernels* 2>/dev/null || ls /opt/rocm*/lib64/rocblas/library/Kernels* 2>/dev/null) | grep -oP 'gfx[0-9a-z+-]+' | sort -u || \
|
||||
echo "WARNING: No rocBLAS kernel data found" \
|
||||
; fi
|
||||
|
||||
RUN echo "TARGETARCH: $TARGETARCH"
|
||||
|
||||
# We need protoc installed, and the version in 22.04 is too old. We will create one as part installing the GRPC build below
|
||||
# but that will also being in a newer version of absl which stablediffusion cannot compile with. This version of protoc is only
|
||||
# here so that we can generate the grpc code for the stablediffusion build
|
||||
RUN <<EOT bash
|
||||
if [ "amd64" = "$TARGETARCH" ]; then
|
||||
curl -L -s https://github.com/protocolbuffers/protobuf/releases/download/v27.1/protoc-27.1-linux-x86_64.zip -o protoc.zip && \
|
||||
unzip -j -d /usr/local/bin protoc.zip bin/protoc && \
|
||||
rm protoc.zip
|
||||
fi
|
||||
if [ "arm64" = "$TARGETARCH" ]; then
|
||||
curl -L -s https://github.com/protocolbuffers/protobuf/releases/download/v27.1/protoc-27.1-linux-aarch_64.zip -o protoc.zip && \
|
||||
unzip -j -d /usr/local/bin protoc.zip bin/protoc && \
|
||||
rm protoc.zip
|
||||
fi
|
||||
EOT
|
||||
|
||||
# Install CMake (the version in 22.04 is too old)
|
||||
RUN <<EOT bash
|
||||
if [ "${CMAKE_FROM_SOURCE}" = "true" ]; then
|
||||
curl -L -s https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/cmake-${CMAKE_VERSION}.tar.gz -o cmake.tar.gz && tar xvf cmake.tar.gz && cd cmake-${CMAKE_VERSION} && ./configure && make && make install
|
||||
else
|
||||
apt-get update && \
|
||||
apt-get install -y \
|
||||
cmake && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
fi
|
||||
EOT
|
||||
|
||||
COPY --from=grpc /opt/grpc /usr/local
|
||||
|
||||
|
||||
COPY . /LocalAI
|
||||
|
||||
RUN <<'EOT' bash
|
||||
set -euxo pipefail
|
||||
|
||||
if [[ -n "${CUDA_DOCKER_ARCH:-}" ]]; then
|
||||
CUDA_ARCH_ESC="${CUDA_DOCKER_ARCH//;/\\;}"
|
||||
export CMAKE_ARGS="${CMAKE_ARGS:-} -DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCH_ESC}"
|
||||
echo "CMAKE_ARGS(env) = ${CMAKE_ARGS}"
|
||||
rm -rf /LocalAI/backend/cpp/turboquant-*-build
|
||||
fi
|
||||
|
||||
cd /LocalAI/backend/cpp/turboquant
|
||||
|
||||
if [ "${TARGETARCH}" = "arm64" ] || [ "${BUILD_TYPE}" = "hipblas" ]; then
|
||||
make turboquant-fallback
|
||||
make turboquant-grpc
|
||||
make turboquant-rpc-server
|
||||
else
|
||||
make turboquant-avx
|
||||
make turboquant-avx2
|
||||
make turboquant-avx512
|
||||
make turboquant-fallback
|
||||
make turboquant-grpc
|
||||
make turboquant-rpc-server
|
||||
fi
|
||||
EOT
|
||||
|
||||
|
||||
# Copy libraries using a script to handle architecture differences
|
||||
RUN make -BC /LocalAI/backend/cpp/turboquant package
|
||||
|
||||
|
||||
FROM scratch
|
||||
|
||||
|
||||
# Copy all available binaries (the build process only creates the appropriate ones for the target architecture)
|
||||
COPY --from=builder /LocalAI/backend/cpp/turboquant/package/. ./
|
||||
@@ -17,12 +17,18 @@ service Backend {
|
||||
rpc GenerateImage(GenerateImageRequest) returns (Result) {}
|
||||
rpc GenerateVideo(GenerateVideoRequest) returns (Result) {}
|
||||
rpc AudioTranscription(TranscriptRequest) returns (TranscriptResult) {}
|
||||
rpc AudioTranscriptionStream(TranscriptRequest) returns (stream TranscriptStreamResponse) {}
|
||||
rpc TTS(TTSRequest) returns (Result) {}
|
||||
rpc TTSStream(TTSRequest) returns (stream Reply) {}
|
||||
rpc SoundGeneration(SoundGenerationRequest) returns (Result) {}
|
||||
rpc TokenizeString(PredictOptions) returns (TokenizationResponse) {}
|
||||
rpc Status(HealthMessage) returns (StatusResponse) {}
|
||||
rpc Detect(DetectOptions) returns (DetectResponse) {}
|
||||
rpc FaceVerify(FaceVerifyRequest) returns (FaceVerifyResponse) {}
|
||||
rpc FaceAnalyze(FaceAnalyzeRequest) returns (FaceAnalyzeResponse) {}
|
||||
rpc VoiceVerify(VoiceVerifyRequest) returns (VoiceVerifyResponse) {}
|
||||
rpc VoiceAnalyze(VoiceAnalyzeRequest) returns (VoiceAnalyzeResponse) {}
|
||||
rpc VoiceEmbed(VoiceEmbedRequest) returns (VoiceEmbedResponse) {}
|
||||
|
||||
rpc StoresSet(StoresSetOptions) returns (Result) {}
|
||||
rpc StoresDelete(StoresDeleteOptions) returns (Result) {}
|
||||
@@ -322,11 +328,21 @@ message TranscriptRequest {
|
||||
bool translate = 5;
|
||||
bool diarize = 6;
|
||||
string prompt = 7;
|
||||
float temperature = 8;
|
||||
repeated string timestamp_granularities = 9;
|
||||
bool stream = 10;
|
||||
}
|
||||
|
||||
message TranscriptResult {
|
||||
repeated TranscriptSegment segments = 1;
|
||||
string text = 2;
|
||||
string language = 3;
|
||||
float duration = 4;
|
||||
}
|
||||
|
||||
message TranscriptStreamResponse {
|
||||
string delta = 1;
|
||||
TranscriptResult final_result = 2;
|
||||
}
|
||||
|
||||
message TranscriptSegment {
|
||||
@@ -464,6 +480,112 @@ message DetectResponse {
|
||||
repeated Detection Detections = 1;
|
||||
}
|
||||
|
||||
// --- Face recognition messages ---
|
||||
|
||||
message FacialArea {
|
||||
float x = 1;
|
||||
float y = 2;
|
||||
float w = 3;
|
||||
float h = 4;
|
||||
}
|
||||
|
||||
message FaceVerifyRequest {
|
||||
string img1 = 1; // base64-encoded image
|
||||
string img2 = 2; // base64-encoded image
|
||||
float threshold = 3; // cosine-distance threshold; 0 = use backend default
|
||||
bool anti_spoofing = 4; // run MiniFASNet liveness on each image; failed liveness forces verified=false
|
||||
}
|
||||
|
||||
message FaceVerifyResponse {
|
||||
bool verified = 1;
|
||||
float distance = 2; // 1 - cosine_similarity
|
||||
float threshold = 3;
|
||||
float confidence = 4; // 0-100
|
||||
string model = 5; // e.g. "buffalo_l"
|
||||
FacialArea img1_area = 6;
|
||||
FacialArea img2_area = 7;
|
||||
float processing_time_ms = 8;
|
||||
bool img1_is_real = 9; // anti-spoofing result when enabled
|
||||
float img1_antispoof_score = 10;
|
||||
bool img2_is_real = 11;
|
||||
float img2_antispoof_score = 12;
|
||||
}
|
||||
|
||||
message FaceAnalyzeRequest {
|
||||
string img = 1; // base64-encoded image
|
||||
repeated string actions = 2; // subset of ["age","gender","emotion","race"]; empty = all-supported
|
||||
bool anti_spoofing = 3;
|
||||
}
|
||||
|
||||
message FaceAnalysis {
|
||||
FacialArea region = 1;
|
||||
float face_confidence = 2;
|
||||
float age = 3;
|
||||
string dominant_gender = 4; // "Man" | "Woman"
|
||||
map<string, float> gender = 5;
|
||||
string dominant_emotion = 6; // reserved; empty in MVP
|
||||
map<string, float> emotion = 7;
|
||||
string dominant_race = 8; // not populated
|
||||
map<string, float> race = 9;
|
||||
bool is_real = 10; // anti-spoofing result when enabled
|
||||
float antispoof_score = 11;
|
||||
}
|
||||
|
||||
message FaceAnalyzeResponse {
|
||||
repeated FaceAnalysis faces = 1;
|
||||
}
|
||||
|
||||
// --- Voice (speaker) recognition messages ---
|
||||
//
|
||||
// Analogous to the Face* messages above, but for speaker biometrics.
|
||||
// Audio fields accept a filesystem path (same convention as
|
||||
// TranscriptRequest.dst). The HTTP layer materialises base64 / URL /
|
||||
// data-URI inputs to a temp file before calling the gRPC backend.
|
||||
|
||||
message VoiceVerifyRequest {
|
||||
string audio1 = 1; // path to first audio clip
|
||||
string audio2 = 2; // path to second audio clip
|
||||
float threshold = 3; // cosine-distance threshold; 0 = use backend default
|
||||
bool anti_spoofing = 4; // reserved for future AASIST bolt-on
|
||||
}
|
||||
|
||||
message VoiceVerifyResponse {
|
||||
bool verified = 1;
|
||||
float distance = 2; // 1 - cosine_similarity
|
||||
float threshold = 3;
|
||||
float confidence = 4; // 0-100
|
||||
string model = 5; // e.g. "speechbrain/spkrec-ecapa-voxceleb"
|
||||
float processing_time_ms = 6;
|
||||
}
|
||||
|
||||
message VoiceAnalyzeRequest {
|
||||
string audio = 1; // path to audio clip
|
||||
repeated string actions = 2; // subset of ["age","gender","emotion"]; empty = all-supported
|
||||
}
|
||||
|
||||
message VoiceAnalysis {
|
||||
float start = 1; // segment start time in seconds (0 if single-utterance)
|
||||
float end = 2; // segment end time in seconds
|
||||
float age = 3;
|
||||
string dominant_gender = 4;
|
||||
map<string, float> gender = 5;
|
||||
string dominant_emotion = 6;
|
||||
map<string, float> emotion = 7;
|
||||
}
|
||||
|
||||
message VoiceAnalyzeResponse {
|
||||
repeated VoiceAnalysis segments = 1;
|
||||
}
|
||||
|
||||
message VoiceEmbedRequest {
|
||||
string audio = 1; // path to audio clip
|
||||
}
|
||||
|
||||
message VoiceEmbedResponse {
|
||||
repeated float embedding = 1;
|
||||
string model = 2;
|
||||
}
|
||||
|
||||
message ToolFormatMarkers {
|
||||
string format_type = 1; // "json_native", "tag_with_json", "tag_with_tagged"
|
||||
|
||||
@@ -546,6 +668,7 @@ message ModelMetadataResponse {
|
||||
bool supports_thinking = 1;
|
||||
string rendered_template = 2; // The rendered chat template with enable_thinking=true (empty if not applicable)
|
||||
ToolFormatMarkers tool_format = 3; // Auto-detected tool format markers from differential template analysis
|
||||
string media_marker = 4; // Marker the backend expects in the prompt for each multimodal input (images/audio/video). Empty when the backend does not use a marker.
|
||||
}
|
||||
|
||||
// Fine-tuning messages
|
||||
|
||||
85
backend/cpp/buun-llama-cpp/Makefile
Normal file
85
backend/cpp/buun-llama-cpp/Makefile
Normal file
@@ -0,0 +1,85 @@
|
||||
|
||||
# Pinned to the HEAD of master on https://github.com/spiritbuun/buun-llama-cpp.
|
||||
# Auto-bumped nightly by .github/workflows/bump_deps.yaml.
|
||||
BUUN_LLAMA_VERSION?=22464d0848b87c5d56b52fdf6af2e5da46bf803e
|
||||
LLAMA_REPO?=https://github.com/spiritbuun/buun-llama-cpp
|
||||
|
||||
CMAKE_ARGS?=
|
||||
BUILD_TYPE?=
|
||||
NATIVE?=false
|
||||
ONEAPI_VARS?=/opt/intel/oneapi/setvars.sh
|
||||
TARGET?=--target grpc-server
|
||||
JOBS?=$(shell nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 1)
|
||||
ARCH?=$(shell uname -m)
|
||||
|
||||
CURRENT_MAKEFILE_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
|
||||
LLAMA_CPP_DIR := $(CURRENT_MAKEFILE_DIR)/../llama-cpp
|
||||
|
||||
GREEN := \033[0;32m
|
||||
RESET := \033[0m
|
||||
|
||||
# buun-llama-cpp is a llama.cpp fork-of-a-fork (spiritbuun/buun-llama-cpp forked
|
||||
# TheTom/llama-cpp-turboquant, which itself forked ggml-org/llama.cpp). Rather
|
||||
# than duplicating grpc-server.cpp / CMakeLists.txt / prepare.sh we reuse the
|
||||
# ones in backend/cpp/llama-cpp, and only swap which repo+sha the fetch step
|
||||
# pulls. Each flavor target copies ../llama-cpp into a sibling
|
||||
# ../buun-llama-cpp-<flavor>-build directory, then invokes llama-cpp's own
|
||||
# build-llama-cpp-grpc-server with LLAMA_REPO/LLAMA_VERSION overridden to point
|
||||
# at the fork.
|
||||
PATCHES_DIR := $(CURRENT_MAKEFILE_DIR)/patches
|
||||
|
||||
# Each flavor target:
|
||||
# 1. copies backend/cpp/llama-cpp/ (grpc-server.cpp + prepare.sh + CMakeLists.txt + Makefile)
|
||||
# into a sibling buun-llama-cpp-<flavor>-build directory;
|
||||
# 2. clones the buun fork into buun-llama-cpp-<flavor>-build/llama.cpp via the
|
||||
# copy's own `llama.cpp` target, overriding LLAMA_REPO/LLAMA_VERSION;
|
||||
# 3. applies patches from backend/cpp/buun-llama-cpp/patches/ to the cloned
|
||||
# fork sources (for backporting upstream commits the fork hasn't pulled);
|
||||
# 4. runs the copy's `grpc-server` target, which produces the binary we copy
|
||||
# up as buun-llama-cpp-<flavor>.
|
||||
define buun-llama-cpp-build
|
||||
rm -rf $(CURRENT_MAKEFILE_DIR)/../buun-llama-cpp-$(1)-build
|
||||
cp -rf $(LLAMA_CPP_DIR) $(CURRENT_MAKEFILE_DIR)/../buun-llama-cpp-$(1)-build
|
||||
$(MAKE) -C $(CURRENT_MAKEFILE_DIR)/../buun-llama-cpp-$(1)-build purge
|
||||
# Augment the copied grpc-server.cpp's KV-cache allow-list with the
|
||||
# fork's turbo2/turbo3/turbo4/turbo2_tcq/turbo3_tcq types and wire up the
|
||||
# DFlash-specific option handlers (tree_budget / draft_topk). We patch the
|
||||
# *copy*, never the original under backend/cpp/llama-cpp/, so the stock
|
||||
# llama-cpp build stays compiling against vanilla upstream.
|
||||
bash $(CURRENT_MAKEFILE_DIR)/patch-grpc-server.sh $(CURRENT_MAKEFILE_DIR)/../buun-llama-cpp-$(1)-build/grpc-server.cpp
|
||||
$(info $(GREEN)I buun-llama-cpp build info:$(1)$(RESET))
|
||||
LLAMA_REPO=$(LLAMA_REPO) LLAMA_VERSION=$(BUUN_LLAMA_VERSION) \
|
||||
$(MAKE) -C $(CURRENT_MAKEFILE_DIR)/../buun-llama-cpp-$(1)-build llama.cpp
|
||||
bash $(CURRENT_MAKEFILE_DIR)/apply-patches.sh $(CURRENT_MAKEFILE_DIR)/../buun-llama-cpp-$(1)-build/llama.cpp $(PATCHES_DIR)
|
||||
CMAKE_ARGS="$(CMAKE_ARGS) $(2)" TARGET="$(3)" \
|
||||
LLAMA_REPO=$(LLAMA_REPO) LLAMA_VERSION=$(BUUN_LLAMA_VERSION) \
|
||||
$(MAKE) -C $(CURRENT_MAKEFILE_DIR)/../buun-llama-cpp-$(1)-build grpc-server
|
||||
cp -rfv $(CURRENT_MAKEFILE_DIR)/../buun-llama-cpp-$(1)-build/grpc-server buun-llama-cpp-$(1)
|
||||
endef
|
||||
|
||||
buun-llama-cpp-avx2:
|
||||
$(call buun-llama-cpp-build,avx2,-DGGML_AVX=on -DGGML_AVX2=on -DGGML_AVX512=off -DGGML_FMA=on -DGGML_F16C=on,--target grpc-server)
|
||||
|
||||
buun-llama-cpp-avx512:
|
||||
$(call buun-llama-cpp-build,avx512,-DGGML_AVX=on -DGGML_AVX2=off -DGGML_AVX512=on -DGGML_FMA=on -DGGML_F16C=on,--target grpc-server)
|
||||
|
||||
buun-llama-cpp-avx:
|
||||
$(call buun-llama-cpp-build,avx,-DGGML_AVX=on -DGGML_AVX2=off -DGGML_AVX512=off -DGGML_FMA=off -DGGML_F16C=off -DGGML_BMI2=off,--target grpc-server)
|
||||
|
||||
buun-llama-cpp-fallback:
|
||||
$(call buun-llama-cpp-build,fallback,-DGGML_AVX=off -DGGML_AVX2=off -DGGML_AVX512=off -DGGML_FMA=off -DGGML_F16C=off -DGGML_BMI2=off,--target grpc-server)
|
||||
|
||||
buun-llama-cpp-grpc:
|
||||
$(call buun-llama-cpp-build,grpc,-DGGML_RPC=ON -DGGML_AVX=off -DGGML_AVX2=off -DGGML_AVX512=off -DGGML_FMA=off -DGGML_F16C=off -DGGML_BMI2=off,--target grpc-server --target rpc-server)
|
||||
|
||||
buun-llama-cpp-rpc-server: buun-llama-cpp-grpc
|
||||
cp -rf $(CURRENT_MAKEFILE_DIR)/../buun-llama-cpp-grpc-build/llama.cpp/build/bin/rpc-server buun-llama-cpp-rpc-server
|
||||
|
||||
package:
|
||||
bash package.sh
|
||||
|
||||
purge:
|
||||
rm -rf $(CURRENT_MAKEFILE_DIR)/../buun-llama-cpp-*-build
|
||||
rm -rf buun-llama-cpp-* package
|
||||
|
||||
clean: purge
|
||||
50
backend/cpp/buun-llama-cpp/apply-patches.sh
Executable file
50
backend/cpp/buun-llama-cpp/apply-patches.sh
Executable file
@@ -0,0 +1,50 @@
|
||||
#!/bin/bash
|
||||
# Apply the buun-llama-cpp patch series to a cloned buun-llama-cpp checkout.
|
||||
#
|
||||
# buun-llama-cpp is a fork-of-a-fork that branched off upstream llama.cpp
|
||||
# before some API changes the shared backend/cpp/llama-cpp/grpc-server.cpp
|
||||
# depends on. We carry those upstream commits as patch files under
|
||||
# backend/cpp/buun-llama-cpp/patches/ and apply them here so the reused
|
||||
# grpc-server source compiles against the fork unmodified.
|
||||
#
|
||||
# Drop the corresponding patch from patches/ whenever the fork catches up with
|
||||
# upstream — the build will fail fast if a patch stops applying, which is the
|
||||
# signal to retire it.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
if [[ $# -ne 2 ]]; then
|
||||
echo "usage: $0 <llama.cpp-src-dir> <patches-dir>" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
SRC_DIR=$1
|
||||
PATCHES_DIR=$2
|
||||
|
||||
if [[ ! -d "$SRC_DIR" ]]; then
|
||||
echo "source dir does not exist: $SRC_DIR" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
if [[ ! -d "$PATCHES_DIR" ]]; then
|
||||
echo "no patches dir at $PATCHES_DIR, nothing to apply"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
shopt -s nullglob
|
||||
patches=("$PATCHES_DIR"/*.patch)
|
||||
shopt -u nullglob
|
||||
|
||||
if [[ ${#patches[@]} -eq 0 ]]; then
|
||||
echo "no .patch files in $PATCHES_DIR, nothing to apply"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
cd "$SRC_DIR"
|
||||
|
||||
for patch in "${patches[@]}"; do
|
||||
echo "==> applying $patch"
|
||||
git apply --verbose "$patch"
|
||||
done
|
||||
|
||||
echo "all buun-llama-cpp patches applied successfully"
|
||||
57
backend/cpp/buun-llama-cpp/package.sh
Executable file
57
backend/cpp/buun-llama-cpp/package.sh
Executable file
@@ -0,0 +1,57 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Script to copy the appropriate libraries based on architecture
|
||||
# This script is used in the final stage of the Dockerfile
|
||||
|
||||
set -e
|
||||
|
||||
CURDIR=$(dirname "$(realpath $0)")
|
||||
REPO_ROOT="${CURDIR}/../../.."
|
||||
|
||||
# Create lib directory
|
||||
mkdir -p $CURDIR/package/lib
|
||||
|
||||
cp -avrf $CURDIR/buun-llama-cpp-* $CURDIR/package/
|
||||
cp -rfv $CURDIR/run.sh $CURDIR/package/
|
||||
|
||||
# Detect architecture and copy appropriate libraries
|
||||
if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then
|
||||
# x86_64 architecture
|
||||
echo "Detected x86_64 architecture, copying x86_64 libraries..."
|
||||
cp -arfLv /lib64/ld-linux-x86-64.so.2 $CURDIR/package/lib/ld.so
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libc.so.6 $CURDIR/package/lib/libc.so.6
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libgcc_s.so.1 $CURDIR/package/lib/libgcc_s.so.1
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libstdc++.so.6 $CURDIR/package/lib/libstdc++.so.6
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libm.so.6 $CURDIR/package/lib/libm.so.6
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libgomp.so.1 $CURDIR/package/lib/libgomp.so.1
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libdl.so.2 $CURDIR/package/lib/libdl.so.2
|
||||
cp -arfLv /lib/x86_64-linux-gnu/librt.so.1 $CURDIR/package/lib/librt.so.1
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libpthread.so.0 $CURDIR/package/lib/libpthread.so.0
|
||||
elif [ -f "/lib/ld-linux-aarch64.so.1" ]; then
|
||||
# ARM64 architecture
|
||||
echo "Detected ARM64 architecture, copying ARM64 libraries..."
|
||||
cp -arfLv /lib/ld-linux-aarch64.so.1 $CURDIR/package/lib/ld.so
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libc.so.6 $CURDIR/package/lib/libc.so.6
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libgcc_s.so.1 $CURDIR/package/lib/libgcc_s.so.1
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libstdc++.so.6 $CURDIR/package/lib/libstdc++.so.6
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libm.so.6 $CURDIR/package/lib/libm.so.6
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libgomp.so.1 $CURDIR/package/lib/libgomp.so.1
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libdl.so.2 $CURDIR/package/lib/libdl.so.2
|
||||
cp -arfLv /lib/aarch64-linux-gnu/librt.so.1 $CURDIR/package/lib/librt.so.1
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libpthread.so.0 $CURDIR/package/lib/libpthread.so.0
|
||||
else
|
||||
echo "Error: Could not detect architecture"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Package GPU libraries based on BUILD_TYPE
|
||||
GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh"
|
||||
if [ -f "$GPU_LIB_SCRIPT" ]; then
|
||||
echo "Packaging GPU libraries for BUILD_TYPE=${BUILD_TYPE:-cpu}..."
|
||||
source "$GPU_LIB_SCRIPT" "$CURDIR/package/lib"
|
||||
package_gpu_libs
|
||||
fi
|
||||
|
||||
echo "Packaging completed successfully"
|
||||
ls -liah $CURDIR/package/
|
||||
ls -liah $CURDIR/package/lib/
|
||||
162
backend/cpp/buun-llama-cpp/patch-grpc-server.sh
Executable file
162
backend/cpp/buun-llama-cpp/patch-grpc-server.sh
Executable file
@@ -0,0 +1,162 @@
|
||||
#!/bin/bash
|
||||
# Patch the shared backend/cpp/llama-cpp/grpc-server.cpp *copy* used by the
|
||||
# buun-llama-cpp build to account for three gaps between upstream and the fork:
|
||||
#
|
||||
# 1. Augment the kv_cache_types[] allow-list so `LoadModel` accepts the
|
||||
# fork-specific `turbo2` / `turbo3` / `turbo4` cache types plus the buun
|
||||
# additions `turbo2_tcq` / `turbo3_tcq`.
|
||||
#
|
||||
# 2. Wire up buun-exclusive speculative-decoding option handlers
|
||||
# (tree_budget / draft_topk) alongside the existing spec_* handlers.
|
||||
# These reference struct fields (common_params.speculative.tree_budget
|
||||
# and .draft_topk) that only exist in buun's common/common.h — adding
|
||||
# them to the shared backend/cpp/llama-cpp/grpc-server.cpp would break
|
||||
# the stock llama-cpp build, so we inject them only into the buun copy.
|
||||
#
|
||||
# 3. Replace `get_media_marker()` (added upstream in ggml-org/llama.cpp#21962,
|
||||
# server-side random per-instance marker) with the legacy "<__media__>"
|
||||
# literal. The fork branched before that PR, so server-common.cpp has no
|
||||
# get_media_marker symbol. The fork's mtmd_default_marker() still returns
|
||||
# "<__media__>", and Go-side tooling falls back to that sentinel when the
|
||||
# backend does not expose media_marker, so substituting the literal keeps
|
||||
# behavior identical on the buun path.
|
||||
#
|
||||
# We patch the *copy* sitting in buun-llama-cpp-<flavor>-build/, never the
|
||||
# original under backend/cpp/llama-cpp/, so the stock llama-cpp build keeps
|
||||
# compiling against vanilla upstream.
|
||||
#
|
||||
# Idempotent: skips each insertion if its marker is already present (so re-runs
|
||||
# of the same build dir don't double-insert).
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
if [[ $# -ne 1 ]]; then
|
||||
echo "usage: $0 <grpc-server.cpp>" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
SRC=$1
|
||||
|
||||
if [[ ! -f "$SRC" ]]; then
|
||||
echo "grpc-server.cpp not found at $SRC" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
if grep -q 'GGML_TYPE_TURBO2_TCQ' "$SRC"; then
|
||||
echo "==> $SRC already has buun cache types, skipping KV allow-list patch"
|
||||
else
|
||||
echo "==> patching $SRC to allow turbo2/turbo3/turbo4/turbo2_tcq/turbo3_tcq KV-cache types"
|
||||
|
||||
# Insert the five TURBO entries right after the first ` GGML_TYPE_Q5_1,`
|
||||
# line (the kv_cache_types[] allow-list). Using awk because the builder
|
||||
# image does not ship python3, and GNU sed's multi-line `a\` quoting is
|
||||
# awkward.
|
||||
awk '
|
||||
/^ GGML_TYPE_Q5_1,$/ && !done {
|
||||
print
|
||||
print " // buun-llama-cpp fork extras — added by patch-grpc-server.sh"
|
||||
print " GGML_TYPE_TURBO2_0,"
|
||||
print " GGML_TYPE_TURBO3_0,"
|
||||
print " GGML_TYPE_TURBO4_0,"
|
||||
print " GGML_TYPE_TURBO2_TCQ,"
|
||||
print " GGML_TYPE_TURBO3_TCQ,"
|
||||
done = 1
|
||||
next
|
||||
}
|
||||
{ print }
|
||||
END {
|
||||
if (!done) {
|
||||
print "patch-grpc-server.sh: anchor ` GGML_TYPE_Q5_1,` not found" > "/dev/stderr"
|
||||
exit 1
|
||||
}
|
||||
}
|
||||
' "$SRC" > "$SRC.tmp"
|
||||
mv "$SRC.tmp" "$SRC"
|
||||
|
||||
echo "==> KV allow-list patch OK"
|
||||
fi
|
||||
|
||||
if grep -q 'optname, "tree_budget"' "$SRC"; then
|
||||
echo "==> $SRC already has DFlash option handlers, skipping"
|
||||
else
|
||||
echo "==> patching $SRC to add tree_budget / draft_topk option handlers"
|
||||
|
||||
# Insert two new `else if` handlers between the inner close-brace of the
|
||||
# `spec_p_split` block and the next `} else if (…spec_ngram_size_n…)` line.
|
||||
# Upstream writes each `} else if` as a single physical line, so we don't
|
||||
# emit an outer `}` ourselves — the existing next line provides both the
|
||||
# close of our `draft_topk` block and the open of `spec_ngram_size_n`.
|
||||
# Anchor on the exact 3-line body of spec_p_split so we can't drift.
|
||||
awk '
|
||||
prev2 == " } else if (!strcmp(optname, \"spec_p_split\")) {" &&
|
||||
prev1 ~ /^ +if \(optval != NULL\) \{$/ &&
|
||||
$0 ~ /^ +try \{ params\.speculative\.p_split = std::stof\(optval_str\); \} catch \(\.\.\.\) \{\}$/ &&
|
||||
!done {
|
||||
print # print the try-line itself
|
||||
getline inner_close # read " }" closing the inner if
|
||||
print inner_close # print it — this closes spec_p_split body
|
||||
print " // buun-llama-cpp DFlash options — added by patch-grpc-server.sh"
|
||||
print " } else if (!strcmp(optname, \"tree_budget\")) {"
|
||||
print " if (optval != NULL) {"
|
||||
print " try { params.speculative.tree_budget = std::stoi(optval_str); } catch (...) {}"
|
||||
print " }"
|
||||
print " } else if (!strcmp(optname, \"draft_topk\")) {"
|
||||
print " if (optval != NULL) {"
|
||||
print " try { params.speculative.draft_topk = std::stoi(optval_str); } catch (...) {}"
|
||||
print " }"
|
||||
# The next source line (`} else if (…spec_ngram_size_n…) {`) closes
|
||||
# our draft_topk block and continues the chain naturally; fall back
|
||||
# into the main loop to emit it and everything after.
|
||||
done = 1
|
||||
prev2 = prev1
|
||||
prev1 = inner_close
|
||||
next
|
||||
}
|
||||
{ print; prev2 = prev1; prev1 = $0 }
|
||||
END {
|
||||
if (!done) {
|
||||
print "patch-grpc-server.sh: spec_p_split anchor not found" > "/dev/stderr"
|
||||
exit 1
|
||||
}
|
||||
}
|
||||
' "$SRC" > "$SRC.tmp"
|
||||
mv "$SRC.tmp" "$SRC"
|
||||
|
||||
echo "==> DFlash option-handler patch OK"
|
||||
fi
|
||||
|
||||
if grep -qE 'ctx_server\.get_meta\(\)\.logit_bias_eog|params_base\.sampling\.logit_bias_eog,' "$SRC"; then
|
||||
echo "==> patching $SRC to drop the logit_bias_eog arg from params_from_json_cmpl() callsites (buun still uses the pre-refactor 4-arg signature)"
|
||||
# Upstream llama.cpp refactored params_from_json_cmpl to take a precomputed
|
||||
# logit_bias_eog vector after buun's 2026-04-05 fork-point — simultaneously
|
||||
# adding server_context_meta::logit_bias_eog as the supplier. Buun carries
|
||||
# neither change: its params_from_json_cmpl is still 4-arg, and internally
|
||||
# derives logit_bias_eog from the common_params it's passed. So we just
|
||||
# delete the argument line entirely — the remaining 4 args match buun's
|
||||
# signature and the resulting behavior matches upstream bit-for-bit
|
||||
# (upstream's 5th arg is the same data buun derives internally).
|
||||
#
|
||||
# Guard is broad so this works whether the line has been run through this
|
||||
# block before (leaving params_base.sampling.logit_bias_eog,) or not
|
||||
# (leaving the original ctx_server.get_meta().logit_bias_eog,).
|
||||
sed -E '/^[[:space:]]+(ctx_server\.get_meta\(\)\.logit_bias_eog|params_base\.sampling\.logit_bias_eog),$/d' "$SRC" > "$SRC.tmp"
|
||||
mv "$SRC.tmp" "$SRC"
|
||||
echo "==> logit_bias_eog arg drop OK"
|
||||
else
|
||||
echo "==> $SRC has no logit_bias_eog arg line, skipping"
|
||||
fi
|
||||
|
||||
if grep -q 'get_media_marker()' "$SRC"; then
|
||||
echo "==> patching $SRC to replace get_media_marker() with legacy \"<__media__>\" literal"
|
||||
# Only one call site today (ModelMetadata), but replace all occurrences to
|
||||
# stay robust if upstream adds more. Use a temp file to avoid relying on
|
||||
# sed -i portability (the builder image uses GNU sed, but keeping this
|
||||
# consistent with the awk block above).
|
||||
sed 's/get_media_marker()/"<__media__>"/g' "$SRC" > "$SRC.tmp"
|
||||
mv "$SRC.tmp" "$SRC"
|
||||
echo "==> get_media_marker() substitution OK"
|
||||
else
|
||||
echo "==> $SRC has no get_media_marker() call, skipping media-marker patch"
|
||||
fi
|
||||
|
||||
echo "==> all patches applied"
|
||||
@@ -0,0 +1,46 @@
|
||||
Subject: [PATCH] ggml-cuda/fattn: provide atomicAdd(double*,double) shim for pre-sm_60
|
||||
|
||||
Buun's Q² calibration path in ggml_cuda_turbo_scale_q calls
|
||||
atomicAdd(&d_q_channel_sq_fattn[threadIdx.x], (double)(val * val));
|
||||
but native double atomicAdd is only available on compute capability 6.0
|
||||
and newer. Compiling against a CUDA arch list that includes older
|
||||
architectures (LocalAI's CUDA 12 Docker image builds for the full
|
||||
published arch range) fails with:
|
||||
|
||||
fattn.cu(812): error: no instance of overloaded function "atomicAdd"
|
||||
matches the argument list, argument types are: (double *, double)
|
||||
|
||||
Add the canonical CUDA-programming-guide shim at the top of fattn.cu so
|
||||
pre-sm_60 codegen has a definition to call. On sm_60+ the native CUDA
|
||||
intrinsic is used and the shim is elided via __CUDA_ARCH__.
|
||||
|
||||
--- a/ggml/src/ggml-cuda/fattn.cu
|
||||
+++ b/ggml/src/ggml-cuda/fattn.cu
|
||||
@@ -7,6 +7,27 @@
|
||||
|
||||
#include <atomic>
|
||||
|
||||
+// Pre-sm_60 double atomicAdd shim. Native double atomicAdd(double*,double)
|
||||
+// is only available on CUDA compute capability 6.0+ (see CUDA C Programming
|
||||
+// Guide, B.15 Atomic Functions). Buun's Q² calibration path below calls
|
||||
+// atomicAdd with a double*; without this definition, nvcc fails to find a
|
||||
+// matching overload whenever the compile target list includes pre-sm_60
|
||||
+// architectures. The standard CAS loop implementation below matches the
|
||||
+// semantics of the native intrinsic.
|
||||
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 600
|
||||
+static __device__ double atomicAdd(double * address, double val) {
|
||||
+ unsigned long long int * address_as_ull = (unsigned long long int *)address;
|
||||
+ unsigned long long int old = *address_as_ull;
|
||||
+ unsigned long long int assumed;
|
||||
+ do {
|
||||
+ assumed = old;
|
||||
+ old = atomicCAS(address_as_ull, assumed,
|
||||
+ __double_as_longlong(val + __longlong_as_double(assumed)));
|
||||
+ } while (assumed != old);
|
||||
+ return __longlong_as_double(old);
|
||||
+}
|
||||
+#endif
|
||||
+
|
||||
// InnerQ: update the fattn-side inverse scale array from host (all devices)
|
||||
void turbo_innerq_update_fattn_scales(const float * scale_inv) {
|
||||
int cur_device;
|
||||
@@ -0,0 +1,32 @@
|
||||
Subject: [PATCH] ggml-cuda/argmax: pass WARP_SIZE to the top-K __shfl_xor_sync calls
|
||||
|
||||
Two __shfl_xor_sync calls in the top-K intra-warp merge drop the `width`
|
||||
argument and rely on the CUDA default (warpSize). Every other call in
|
||||
the same file already passes WARP_SIZE explicitly, and the HIP/ROCm
|
||||
compatibility shim at ggml/src/ggml-cuda/vendors/hip.h:33 is a 4-arg
|
||||
function-like macro — so the 3-arg form fails to preprocess when
|
||||
building with hipcc against ROCm:
|
||||
|
||||
argmax.cu:265: error: too few arguments provided to function-like
|
||||
macro invocation
|
||||
note: macro '__shfl_xor_sync' defined here:
|
||||
#define __shfl_xor_sync(mask, var, laneMask, width) \
|
||||
__shfl_xor(var, laneMask, width)
|
||||
|
||||
Align the two call sites with the rest of the file by passing WARP_SIZE
|
||||
explicitly. On CUDA the generated code is unchanged (warpSize is the
|
||||
default); on HIP it now matches the macro's arity.
|
||||
|
||||
--- a/ggml/src/ggml-cuda/argmax.cu
|
||||
+++ b/ggml/src/ggml-cuda/argmax.cu
|
||||
@@ -262,8 +262,8 @@
|
||||
// Each step: lane gets partner's min element, if it beats our min, replace and re-heapify
|
||||
for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
|
||||
for (int i = 0; i < K; i++) {
|
||||
- float partner_val = __shfl_xor_sync(0xFFFFFFFF, heap_val[i], offset);
|
||||
- int partner_idx = __shfl_xor_sync(0xFFFFFFFF, heap_idx[i], offset);
|
||||
+ float partner_val = __shfl_xor_sync(0xFFFFFFFF, heap_val[i], offset, WARP_SIZE);
|
||||
+ int partner_idx = __shfl_xor_sync(0xFFFFFFFF, heap_idx[i], offset, WARP_SIZE);
|
||||
if (partner_val > heap_val[0]) {
|
||||
heap_val[0] = partner_val;
|
||||
heap_idx[0] = partner_idx;
|
||||
@@ -0,0 +1,24 @@
|
||||
Subject: [PATCH] ggml-cuda/vendors/hip: alias cudaMemcpy{To,From}Symbol to hip counterparts
|
||||
|
||||
Buun's Q² calibration + TCQ codebook upload paths in fattn.cu use
|
||||
cudaMemcpyToSymbol / cudaMemcpyFromSymbol. The HIP-compat header in
|
||||
ggml/src/ggml-cuda/vendors/hip.h already aliases the scalar cudaMemcpy
|
||||
family (cudaMemcpy, cudaMemcpyAsync, cudaMemcpy2DAsync, …) but is
|
||||
missing the symbol variants. Building with hipcc therefore fails with
|
||||
15+ "use of undeclared identifier 'cudaMemcpyToSymbol'" errors.
|
||||
|
||||
Add the two missing aliases alongside the existing memcpy block. HIP
|
||||
provides hipMemcpy{To,From}Symbol with the same signature as CUDA's
|
||||
equivalents, so this is a straight name substitution.
|
||||
|
||||
--- a/ggml/src/ggml-cuda/vendors/hip.h
|
||||
+++ b/ggml/src/ggml-cuda/vendors/hip.h
|
||||
@@ -85,6 +85,8 @@
|
||||
#define cudaMemcpyDeviceToDevice hipMemcpyDeviceToDevice
|
||||
#define cudaMemcpyDeviceToHost hipMemcpyDeviceToHost
|
||||
#define cudaMemcpyHostToDevice hipMemcpyHostToDevice
|
||||
+#define cudaMemcpyToSymbol hipMemcpyToSymbol
|
||||
+#define cudaMemcpyFromSymbol hipMemcpyFromSymbol
|
||||
#define cudaMemcpyKind hipMemcpyKind
|
||||
#define cudaMemset hipMemset
|
||||
#define cudaMemsetAsync hipMemsetAsync
|
||||
@@ -0,0 +1,36 @@
|
||||
Subject: [PATCH] ggml-cuda/fattn: pass WARP_SIZE to fwht128 __shfl_xor_sync calls
|
||||
|
||||
Same issue as the argmax top-K fix: two __shfl_xor_sync call sites in
|
||||
the FWHT-128 butterfly kernels (ggml_cuda_fwht128 and fwht128_store_half)
|
||||
use the 3-arg CUDA form and omit the `width` argument that the HIP
|
||||
function-like macro in vendors/hip.h:33 requires. Hipcc fails with:
|
||||
|
||||
fattn.cu:512: too few arguments provided to function-like macro
|
||||
invocation
|
||||
note: macro '__shfl_xor_sync' defined here:
|
||||
#define __shfl_xor_sync(mask, var, laneMask, width) \
|
||||
__shfl_xor(var, laneMask, width)
|
||||
|
||||
Add WARP_SIZE to both calls. CUDA codegen is unchanged (warpSize is the
|
||||
default); HIP now matches the macro arity.
|
||||
|
||||
--- a/ggml/src/ggml-cuda/fattn.cu
|
||||
+++ b/ggml/src/ggml-cuda/fattn.cu
|
||||
@@ -509,7 +509,7 @@
|
||||
// Intra-warp passes: shuffle xor with stride h, no smem, no sync.
|
||||
#pragma unroll
|
||||
for (int h = 1; h <= 16; h *= 2) {
|
||||
- const float other = __shfl_xor_sync(0xFFFFFFFF, val, h);
|
||||
+ const float other = __shfl_xor_sync(0xFFFFFFFF, val, h, WARP_SIZE);
|
||||
val = (tid & h) ? (other - val) : (val + other);
|
||||
}
|
||||
|
||||
@@ -533,7 +533,7 @@
|
||||
static __device__ __forceinline__ void fwht128_store_half(
|
||||
float val, half * dst_base) {
|
||||
const int tid = threadIdx.x;
|
||||
- const float neighbor = __shfl_xor_sync(0xFFFFFFFF, val, 1);
|
||||
+ const float neighbor = __shfl_xor_sync(0xFFFFFFFF, val, 1, WARP_SIZE);
|
||||
if ((tid & 1) == 0) {
|
||||
const half2 packed = __floats2half2_rn(val, neighbor);
|
||||
*((half2 *)(dst_base + tid)) = packed;
|
||||
65
backend/cpp/buun-llama-cpp/run.sh
Executable file
65
backend/cpp/buun-llama-cpp/run.sh
Executable file
@@ -0,0 +1,65 @@
|
||||
#!/bin/bash
|
||||
set -ex
|
||||
|
||||
# Get the absolute current dir where the script is located
|
||||
CURDIR=$(dirname "$(realpath $0)")
|
||||
|
||||
cd /
|
||||
|
||||
echo "CPU info:"
|
||||
grep -e "model\sname" /proc/cpuinfo | head -1
|
||||
grep -e "flags" /proc/cpuinfo | head -1
|
||||
|
||||
BINARY=buun-llama-cpp-fallback
|
||||
|
||||
if grep -q -e "\savx\s" /proc/cpuinfo ; then
|
||||
echo "CPU: AVX found OK"
|
||||
if [ -e $CURDIR/buun-llama-cpp-avx ]; then
|
||||
BINARY=buun-llama-cpp-avx
|
||||
fi
|
||||
fi
|
||||
|
||||
if grep -q -e "\savx2\s" /proc/cpuinfo ; then
|
||||
echo "CPU: AVX2 found OK"
|
||||
if [ -e $CURDIR/buun-llama-cpp-avx2 ]; then
|
||||
BINARY=buun-llama-cpp-avx2
|
||||
fi
|
||||
fi
|
||||
|
||||
# Check avx 512
|
||||
if grep -q -e "\savx512f\s" /proc/cpuinfo ; then
|
||||
echo "CPU: AVX512F found OK"
|
||||
if [ -e $CURDIR/buun-llama-cpp-avx512 ]; then
|
||||
BINARY=buun-llama-cpp-avx512
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ -n "$LLAMACPP_GRPC_SERVERS" ]; then
|
||||
if [ -e $CURDIR/buun-llama-cpp-grpc ]; then
|
||||
BINARY=buun-llama-cpp-grpc
|
||||
fi
|
||||
fi
|
||||
|
||||
# Extend ld library path with the dir where this script is located/lib
|
||||
if [ "$(uname)" == "Darwin" ]; then
|
||||
export DYLD_LIBRARY_PATH=$CURDIR/lib:$DYLD_LIBRARY_PATH
|
||||
else
|
||||
export LD_LIBRARY_PATH=$CURDIR/lib:$LD_LIBRARY_PATH
|
||||
# Tell rocBLAS where to find TensileLibrary data (GPU kernel tuning files)
|
||||
if [ -d "$CURDIR/lib/rocblas/library" ]; then
|
||||
export ROCBLAS_TENSILE_LIBPATH=$CURDIR/lib/rocblas/library
|
||||
fi
|
||||
fi
|
||||
|
||||
# If there is a lib/ld.so, use it
|
||||
if [ -f $CURDIR/lib/ld.so ]; then
|
||||
echo "Using lib/ld.so"
|
||||
echo "Using binary: $BINARY"
|
||||
exec $CURDIR/lib/ld.so $CURDIR/$BINARY "$@"
|
||||
fi
|
||||
|
||||
echo "Using binary: $BINARY"
|
||||
exec $CURDIR/$BINARY "$@"
|
||||
|
||||
# We should never reach this point, however just in case we do, run fallback
|
||||
exec $CURDIR/buun-llama-cpp-fallback "$@"
|
||||
@@ -1,5 +1,5 @@
|
||||
|
||||
IK_LLAMA_VERSION?=08ae48c667e3dcd3025821a8585190b4a46c2f7c
|
||||
IK_LLAMA_VERSION?=16996aeab772c69b6473597038b2ef0b85297e8b
|
||||
LLAMA_REPO?=https://github.com/ikawrakow/ik_llama.cpp
|
||||
|
||||
CMAKE_ARGS?=
|
||||
|
||||
@@ -326,7 +326,7 @@ struct llama_client_slot
|
||||
char buffer[512];
|
||||
double t_token = t_prompt_processing / num_prompt_tokens_processed;
|
||||
double n_tokens_second = 1e3 / t_prompt_processing * num_prompt_tokens_processed;
|
||||
sprintf(buffer, "prompt eval time = %10.2f ms / %5d tokens (%8.2f ms per token, %8.2f tokens per second)",
|
||||
snprintf(buffer, sizeof(buffer), "prompt eval time = %10.2f ms / %5d tokens (%8.2f ms per token, %8.2f tokens per second)",
|
||||
t_prompt_processing, num_prompt_tokens_processed,
|
||||
t_token, n_tokens_second);
|
||||
LOG_INFO(buffer, {
|
||||
@@ -340,7 +340,7 @@ struct llama_client_slot
|
||||
|
||||
t_token = t_token_generation / n_decoded;
|
||||
n_tokens_second = 1e3 / t_token_generation * n_decoded;
|
||||
sprintf(buffer, "generation eval time = %10.2f ms / %5d runs (%8.2f ms per token, %8.2f tokens per second)",
|
||||
snprintf(buffer, sizeof(buffer), "generation eval time = %10.2f ms / %5d runs (%8.2f ms per token, %8.2f tokens per second)",
|
||||
t_token_generation, n_decoded,
|
||||
t_token, n_tokens_second);
|
||||
LOG_INFO(buffer, {
|
||||
@@ -352,7 +352,7 @@ struct llama_client_slot
|
||||
{"n_tokens_second", n_tokens_second},
|
||||
});
|
||||
|
||||
sprintf(buffer, " total time = %10.2f ms", t_prompt_processing + t_token_generation);
|
||||
snprintf(buffer, sizeof(buffer), " total time = %10.2f ms", t_prompt_processing + t_token_generation);
|
||||
LOG_INFO(buffer, {
|
||||
{"slot_id", id},
|
||||
{"task_id", task_id},
|
||||
@@ -686,7 +686,16 @@ struct llama_server_context
|
||||
slot->sparams.mirostat_eta = json_value(data, "mirostat_eta", default_sparams.mirostat_eta);
|
||||
slot->params.n_keep = json_value(data, "n_keep", slot->params.n_keep);
|
||||
slot->sparams.seed = json_value(data, "seed", default_sparams.seed);
|
||||
slot->sparams.grammar = json_value(data, "grammar", default_sparams.grammar);
|
||||
{
|
||||
// upstream changed common_params_sampling::grammar from std::string to
|
||||
// the common_grammar struct (type + grammar). The incoming JSON still
|
||||
// carries a plain string, so build the user-provided grammar here and
|
||||
// fall back to the server default when the request omits it.
|
||||
std::string grammar_str = json_value(data, "grammar", std::string());
|
||||
slot->sparams.grammar = grammar_str.empty()
|
||||
? default_sparams.grammar
|
||||
: common_grammar{COMMON_GRAMMAR_TYPE_USER, std::move(grammar_str)};
|
||||
}
|
||||
slot->sparams.n_probs = json_value(data, "n_probs", default_sparams.n_probs);
|
||||
slot->sparams.min_keep = json_value(data, "min_keep", default_sparams.min_keep);
|
||||
slot->sparams.grammar_triggers = grammar_triggers;
|
||||
@@ -1232,7 +1241,7 @@ struct llama_server_context
|
||||
// {"logit_bias", slot.sparams.logit_bias},
|
||||
{"n_probs", slot.sparams.n_probs},
|
||||
{"min_keep", slot.sparams.min_keep},
|
||||
{"grammar", slot.sparams.grammar},
|
||||
{"grammar", slot.sparams.grammar.grammar},
|
||||
{"samplers", samplers}
|
||||
};
|
||||
}
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
--- a/examples/llava/clip.cpp
|
||||
+++ b/examples/llava/clip.cpp
|
||||
@@ -2494,7 +2494,7 @@
|
||||
}
|
||||
new_data = work.data();
|
||||
|
||||
- new_size = ggml_quantize_chunk(new_type, f32_data, new_data, 0, n_elms/cur->ne[0], cur->ne[0], nullptr);
|
||||
+ new_size = ggml_quantize_chunk(new_type, f32_data, new_data, 0, n_elms/cur->ne[0], cur->ne[0], nullptr, nullptr);
|
||||
} else {
|
||||
new_type = cur->type;
|
||||
new_data = cur->data;
|
||||
@@ -62,7 +62,18 @@ add_executable(${TARGET} grpc-server.cpp json.hpp httplib.h)
|
||||
target_include_directories(${TARGET} PRIVATE ../llava)
|
||||
target_include_directories(${TARGET} PRIVATE ${CMAKE_SOURCE_DIR})
|
||||
|
||||
target_link_libraries(${TARGET} PRIVATE common llama mtmd ${CMAKE_THREAD_LIBS_INIT} absl::flags hw_grpc_proto
|
||||
# Upstream llama.cpp renamed the `common` helpers library to `llama-common`.
|
||||
# Forks that branched before the rename (e.g. llama-cpp-turboquant) still
|
||||
# expose it as `common`. Detect which one is present so the same CMakeLists
|
||||
# drives both builds — otherwise an unresolved name silently degrades to a
|
||||
# plain `-l` flag and the PUBLIC include dir (where common.h lives) is lost.
|
||||
if (TARGET llama-common)
|
||||
set(_LLAMA_COMMON_TARGET llama-common)
|
||||
else()
|
||||
set(_LLAMA_COMMON_TARGET common)
|
||||
endif()
|
||||
|
||||
target_link_libraries(${TARGET} PRIVATE ${_LLAMA_COMMON_TARGET} llama mtmd ${CMAKE_THREAD_LIBS_INIT} absl::flags hw_grpc_proto
|
||||
absl::flags_parse
|
||||
gRPC::${_REFLECTION}
|
||||
gRPC::${_GRPC_GRPCPP}
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
|
||||
LLAMA_VERSION?=ff5ef8278615a2462b79b50abdf3cc95cfb31c6f
|
||||
LLAMA_VERSION?=187a45637054881ecacf17f8e2f6f8f2ba7df1c7
|
||||
LLAMA_REPO?=https://github.com/ggerganov/llama.cpp
|
||||
|
||||
CMAKE_ARGS?=
|
||||
@@ -33,7 +33,7 @@ else ifeq ($(BUILD_TYPE),hipblas)
|
||||
ROCM_PATH ?= /opt/rocm
|
||||
export CXX=$(ROCM_HOME)/llvm/bin/clang++
|
||||
export CC=$(ROCM_HOME)/llvm/bin/clang
|
||||
AMDGPU_TARGETS?=gfx908,gfx90a,gfx942,gfx950,gfx1030,gfx1100,gfx1101,gfx1102,gfx1200,gfx1201
|
||||
AMDGPU_TARGETS?=gfx908,gfx90a,gfx942,gfx950,gfx1030,gfx1100,gfx1101,gfx1102,gfx1151,gfx1200,gfx1201
|
||||
CMAKE_ARGS+=-DGGML_HIP=ON -DAMDGPU_TARGETS=$(AMDGPU_TARGETS)
|
||||
else ifeq ($(BUILD_TYPE),vulkan)
|
||||
CMAKE_ARGS+=-DGGML_VULKAN=1
|
||||
@@ -132,7 +132,7 @@ llama.cpp:
|
||||
cd llama.cpp && \
|
||||
git init && \
|
||||
git remote add origin $(LLAMA_REPO) && \
|
||||
git fetch origin && \
|
||||
git fetch --all --tags && \
|
||||
git checkout -b build $(LLAMA_VERSION) && \
|
||||
git submodule update --init --recursive --depth 1 --single-branch
|
||||
|
||||
|
||||
@@ -10,6 +10,14 @@
|
||||
#include "server-task.cpp"
|
||||
#include "server-queue.cpp"
|
||||
#include "server-common.cpp"
|
||||
// server-chat.cpp exists only in llama.cpp after the upstream refactor that
|
||||
// split OAI/Anthropic/Responses/transcription conversion helpers out of
|
||||
// server-common.cpp. When present, server-context.cpp and server-task.cpp
|
||||
// above call into it, so we must pull its definitions into this TU or the
|
||||
// link fails. __has_include keeps the source compatible with older pins.
|
||||
#if __has_include("server-chat.cpp")
|
||||
#include "server-chat.cpp"
|
||||
#endif
|
||||
#include "server-context.cpp"
|
||||
|
||||
// LocalAI
|
||||
@@ -26,6 +34,8 @@
|
||||
#include <regex>
|
||||
#include <atomic>
|
||||
#include <cstdlib>
|
||||
#include <fstream>
|
||||
#include <iterator>
|
||||
#include <mutex>
|
||||
#include <signal.h>
|
||||
#include <thread>
|
||||
@@ -76,6 +86,27 @@ static grpc::Status checkAuth(grpc::ServerContext* context) {
|
||||
return grpc::Status(grpc::StatusCode::UNAUTHENTICATED, "invalid token");
|
||||
}
|
||||
|
||||
// Minimal base64 encoder. The C++ backend already pulls in base64_decode from
|
||||
// llama.cpp's server-common.cpp, but no encoder is exposed — and we need one to
|
||||
// hand audio bytes to the existing PredictOptions.audios path (which expects
|
||||
// base64-encoded strings, just like images).
|
||||
static std::string base64_encode_bytes(const unsigned char* data, size_t len) {
|
||||
static const char tbl[] =
|
||||
"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
|
||||
std::string out;
|
||||
out.reserve(((len + 2) / 3) * 4);
|
||||
for (size_t i = 0; i < len; i += 3) {
|
||||
uint32_t triple = (uint32_t(data[i]) << 16);
|
||||
if (i + 1 < len) triple |= (uint32_t(data[i + 1]) << 8);
|
||||
if (i + 2 < len) triple |= uint32_t(data[i + 2]);
|
||||
out.push_back(tbl[(triple >> 18) & 0x3F]);
|
||||
out.push_back(tbl[(triple >> 12) & 0x3F]);
|
||||
out.push_back(i + 1 < len ? tbl[(triple >> 6) & 0x3F] : '=');
|
||||
out.push_back(i + 2 < len ? tbl[triple & 0x3F] : '=');
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
// END LocalAI
|
||||
|
||||
|
||||
@@ -2791,6 +2822,13 @@ public:
|
||||
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
|
||||
}
|
||||
|
||||
// Report the active multimodal media marker so the Go layer can emit the
|
||||
// same string when rendering prompts outside the tokenizer-template path.
|
||||
// Only meaningful when an mtmd context was initialized (vision/audio models).
|
||||
if (ctx_server.impl->mctx != nullptr) {
|
||||
response->set_media_marker(get_media_marker());
|
||||
}
|
||||
|
||||
// Check if chat templates are initialized
|
||||
if (ctx_server.impl->chat_params.tmpls == nullptr) {
|
||||
// If templates are not initialized, we can't detect thinking support
|
||||
@@ -2931,6 +2969,119 @@ public:
|
||||
|
||||
return grpc::Status::OK;
|
||||
}
|
||||
|
||||
// runTranscriptionAsCompletion implements OAI /v1/audio/transcriptions on
|
||||
// top of the existing chat-completion + multimodal-audio pipeline, exactly
|
||||
// the way upstream llama.cpp's server does it (see
|
||||
// tools/server/server-context.cpp post_transcriptions_oai → forwards into
|
||||
// handle_completions_impl with a single user message attaching the audio
|
||||
// file via the mtmd marker).
|
||||
//
|
||||
// We synthesize a backend::PredictOptions with one user message
|
||||
// ("Transcribe audio to text" + optional language hint) and the audio
|
||||
// bytes attached via the existing PredictOptions.audios field, then
|
||||
// delegate to our own Predict() handler. This keeps every multimodal
|
||||
// codepath identical to the chat path and avoids duplicating ~700 lines
|
||||
// of task-construction logic.
|
||||
grpc::Status runTranscriptionAsCompletion(grpc::ServerContext* context,
|
||||
const backend::TranscriptRequest* request,
|
||||
backend::Reply* out_reply) {
|
||||
if (params_base.model.path.empty()) {
|
||||
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
|
||||
}
|
||||
if (request->dst().empty()) {
|
||||
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "dst (audio file path) is required");
|
||||
}
|
||||
|
||||
// Read audio bytes from the path LocalAI's HTTP layer wrote.
|
||||
std::ifstream f(request->dst(), std::ios::binary);
|
||||
if (!f.is_open()) {
|
||||
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "failed to open audio file: " + request->dst());
|
||||
}
|
||||
std::vector<unsigned char> bytes((std::istreambuf_iterator<char>(f)),
|
||||
std::istreambuf_iterator<char>());
|
||||
f.close();
|
||||
if (bytes.empty()) {
|
||||
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "audio file is empty: " + request->dst());
|
||||
}
|
||||
|
||||
std::string b64 = base64_encode_bytes(bytes.data(), bytes.size());
|
||||
|
||||
// Build the same prompt upstream uses in convert_transcriptions_to_chatcmpl.
|
||||
std::string user_prompt = "Transcribe audio to text";
|
||||
if (!request->language().empty()) {
|
||||
user_prompt += " (language: " + request->language() + ")";
|
||||
}
|
||||
if (!request->prompt().empty()) {
|
||||
// Optional context hint from the caller.
|
||||
user_prompt += "\n" + request->prompt();
|
||||
}
|
||||
|
||||
backend::PredictOptions synthetic;
|
||||
synthetic.set_usetokenizertemplate(true);
|
||||
synthetic.set_temperature(request->temperature());
|
||||
// Generation length: leave at 0 so parse_options uses -1 (model default).
|
||||
// The model's stop tokens / EOS handle termination naturally for ASR.
|
||||
backend::Message* msg = synthetic.add_messages();
|
||||
msg->set_role("user");
|
||||
msg->set_content(user_prompt);
|
||||
synthetic.add_audios(b64);
|
||||
|
||||
return Predict(context, &synthetic, out_reply);
|
||||
}
|
||||
|
||||
grpc::Status AudioTranscription(ServerContext* context,
|
||||
const backend::TranscriptRequest* request,
|
||||
backend::TranscriptResult* response) override {
|
||||
auto auth = checkAuth(context);
|
||||
if (!auth.ok()) return auth;
|
||||
|
||||
backend::Reply reply;
|
||||
grpc::Status st = runTranscriptionAsCompletion(context, request, &reply);
|
||||
if (!st.ok()) {
|
||||
return st;
|
||||
}
|
||||
response->set_text(reply.message());
|
||||
if (!request->language().empty()) {
|
||||
response->set_language(request->language());
|
||||
}
|
||||
return grpc::Status::OK;
|
||||
}
|
||||
|
||||
grpc::Status AudioTranscriptionStream(ServerContext* context,
|
||||
const backend::TranscriptRequest* request,
|
||||
grpc::ServerWriter<backend::TranscriptStreamResponse>* writer) override {
|
||||
auto auth = checkAuth(context);
|
||||
if (!auth.ok()) return auth;
|
||||
|
||||
// Buffered streaming: run the transcription as a normal chat
|
||||
// completion, then emit one delta + one final event. Real
|
||||
// token-by-token streaming would require refactoring PredictStream's
|
||||
// 700-line writer-coupled body; the HTTP/SSE contract is identical
|
||||
// either way, and clients that only consume the assembled text don't
|
||||
// notice the difference.
|
||||
backend::Reply reply;
|
||||
grpc::Status st = runTranscriptionAsCompletion(context, request, &reply);
|
||||
if (!st.ok()) {
|
||||
return st;
|
||||
}
|
||||
|
||||
const std::string& text = reply.message();
|
||||
if (!text.empty()) {
|
||||
backend::TranscriptStreamResponse delta_chunk;
|
||||
delta_chunk.set_delta(text);
|
||||
writer->Write(delta_chunk);
|
||||
}
|
||||
|
||||
backend::TranscriptStreamResponse final_chunk;
|
||||
backend::TranscriptResult* final_result = final_chunk.mutable_final_result();
|
||||
final_result->set_text(text);
|
||||
if (!request->language().empty()) {
|
||||
final_result->set_language(request->language());
|
||||
}
|
||||
writer->Write(final_chunk);
|
||||
return grpc::Status::OK;
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
|
||||
81
backend/cpp/turboquant/Makefile
Normal file
81
backend/cpp/turboquant/Makefile
Normal file
@@ -0,0 +1,81 @@
|
||||
|
||||
# Pinned to the HEAD of feature/turboquant-kv-cache on https://github.com/TheTom/llama-cpp-turboquant.
|
||||
# Auto-bumped nightly by .github/workflows/bump_deps.yaml.
|
||||
TURBOQUANT_VERSION?=627ebbc6e27727bd4f65422d8aa60b13404993c8
|
||||
LLAMA_REPO?=https://github.com/TheTom/llama-cpp-turboquant
|
||||
|
||||
CMAKE_ARGS?=
|
||||
BUILD_TYPE?=
|
||||
NATIVE?=false
|
||||
ONEAPI_VARS?=/opt/intel/oneapi/setvars.sh
|
||||
TARGET?=--target grpc-server
|
||||
JOBS?=$(shell nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 1)
|
||||
ARCH?=$(shell uname -m)
|
||||
|
||||
CURRENT_MAKEFILE_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
|
||||
LLAMA_CPP_DIR := $(CURRENT_MAKEFILE_DIR)/../llama-cpp
|
||||
|
||||
GREEN := \033[0;32m
|
||||
RESET := \033[0m
|
||||
|
||||
# turboquant is a llama.cpp fork. Rather than duplicating grpc-server.cpp / CMakeLists.txt /
|
||||
# prepare.sh we reuse the ones in backend/cpp/llama-cpp, and only swap which repo+sha the
|
||||
# fetch step pulls. Each flavor target copies ../llama-cpp into a sibling ../turboquant-<flavor>-build
|
||||
# directory, then invokes llama-cpp's own build-llama-cpp-grpc-server with LLAMA_REPO/LLAMA_VERSION
|
||||
# overridden to point at the fork.
|
||||
PATCHES_DIR := $(CURRENT_MAKEFILE_DIR)/patches
|
||||
|
||||
# Each flavor target:
|
||||
# 1. copies backend/cpp/llama-cpp/ (grpc-server.cpp + prepare.sh + CMakeLists.txt + Makefile)
|
||||
# into a sibling turboquant-<flavor>-build directory;
|
||||
# 2. clones the turboquant fork into turboquant-<flavor>-build/llama.cpp via the copy's
|
||||
# own `llama.cpp` target, overriding LLAMA_REPO/LLAMA_VERSION;
|
||||
# 3. applies patches from backend/cpp/turboquant/patches/ to the cloned fork sources
|
||||
# (needed until the fork catches up with upstream server-context.cpp changes);
|
||||
# 4. runs the copy's `grpc-server` target, which produces the binary we copy up as
|
||||
# turboquant-<flavor>.
|
||||
define turboquant-build
|
||||
rm -rf $(CURRENT_MAKEFILE_DIR)/../turboquant-$(1)-build
|
||||
cp -rf $(LLAMA_CPP_DIR) $(CURRENT_MAKEFILE_DIR)/../turboquant-$(1)-build
|
||||
$(MAKE) -C $(CURRENT_MAKEFILE_DIR)/../turboquant-$(1)-build purge
|
||||
# Augment the copied grpc-server.cpp's KV-cache allow-list with the
|
||||
# fork's turbo2/turbo3/turbo4 types. We patch the *copy*, never the
|
||||
# original under backend/cpp/llama-cpp/, so the stock llama-cpp build
|
||||
# stays compiling against vanilla upstream.
|
||||
bash $(CURRENT_MAKEFILE_DIR)/patch-grpc-server.sh $(CURRENT_MAKEFILE_DIR)/../turboquant-$(1)-build/grpc-server.cpp
|
||||
$(info $(GREEN)I turboquant build info:$(1)$(RESET))
|
||||
LLAMA_REPO=$(LLAMA_REPO) LLAMA_VERSION=$(TURBOQUANT_VERSION) \
|
||||
$(MAKE) -C $(CURRENT_MAKEFILE_DIR)/../turboquant-$(1)-build llama.cpp
|
||||
bash $(CURRENT_MAKEFILE_DIR)/apply-patches.sh $(CURRENT_MAKEFILE_DIR)/../turboquant-$(1)-build/llama.cpp $(PATCHES_DIR)
|
||||
CMAKE_ARGS="$(CMAKE_ARGS) $(2)" TARGET="$(3)" \
|
||||
LLAMA_REPO=$(LLAMA_REPO) LLAMA_VERSION=$(TURBOQUANT_VERSION) \
|
||||
$(MAKE) -C $(CURRENT_MAKEFILE_DIR)/../turboquant-$(1)-build grpc-server
|
||||
cp -rfv $(CURRENT_MAKEFILE_DIR)/../turboquant-$(1)-build/grpc-server turboquant-$(1)
|
||||
endef
|
||||
|
||||
turboquant-avx2:
|
||||
$(call turboquant-build,avx2,-DGGML_AVX=on -DGGML_AVX2=on -DGGML_AVX512=off -DGGML_FMA=on -DGGML_F16C=on,--target grpc-server)
|
||||
|
||||
turboquant-avx512:
|
||||
$(call turboquant-build,avx512,-DGGML_AVX=on -DGGML_AVX2=off -DGGML_AVX512=on -DGGML_FMA=on -DGGML_F16C=on,--target grpc-server)
|
||||
|
||||
turboquant-avx:
|
||||
$(call turboquant-build,avx,-DGGML_AVX=on -DGGML_AVX2=off -DGGML_AVX512=off -DGGML_FMA=off -DGGML_F16C=off -DGGML_BMI2=off,--target grpc-server)
|
||||
|
||||
turboquant-fallback:
|
||||
$(call turboquant-build,fallback,-DGGML_AVX=off -DGGML_AVX2=off -DGGML_AVX512=off -DGGML_FMA=off -DGGML_F16C=off -DGGML_BMI2=off,--target grpc-server)
|
||||
|
||||
turboquant-grpc:
|
||||
$(call turboquant-build,grpc,-DGGML_RPC=ON -DGGML_AVX=off -DGGML_AVX2=off -DGGML_AVX512=off -DGGML_FMA=off -DGGML_F16C=off -DGGML_BMI2=off,--target grpc-server --target rpc-server)
|
||||
|
||||
turboquant-rpc-server: turboquant-grpc
|
||||
cp -rf $(CURRENT_MAKEFILE_DIR)/../turboquant-grpc-build/llama.cpp/build/bin/rpc-server turboquant-rpc-server
|
||||
|
||||
package:
|
||||
bash package.sh
|
||||
|
||||
purge:
|
||||
rm -rf $(CURRENT_MAKEFILE_DIR)/../turboquant-*-build
|
||||
rm -rf turboquant-* package
|
||||
|
||||
clean: purge
|
||||
50
backend/cpp/turboquant/apply-patches.sh
Executable file
50
backend/cpp/turboquant/apply-patches.sh
Executable file
@@ -0,0 +1,50 @@
|
||||
#!/bin/bash
|
||||
# Apply the turboquant patch series to a cloned llama-cpp-turboquant checkout.
|
||||
#
|
||||
# The turboquant fork branched from upstream llama.cpp before a few API changes
|
||||
# that the shared backend/cpp/llama-cpp/grpc-server.cpp depends on. We carry
|
||||
# those upstream commits as patch files under backend/cpp/turboquant/patches/
|
||||
# and apply them here so the reused grpc-server source compiles against the
|
||||
# fork unmodified.
|
||||
#
|
||||
# Drop the corresponding patch from patches/ whenever the fork catches up with
|
||||
# upstream — the build will fail fast if a patch stops applying, which is the
|
||||
# signal to retire it.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
if [[ $# -ne 2 ]]; then
|
||||
echo "usage: $0 <llama.cpp-src-dir> <patches-dir>" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
SRC_DIR=$1
|
||||
PATCHES_DIR=$2
|
||||
|
||||
if [[ ! -d "$SRC_DIR" ]]; then
|
||||
echo "source dir does not exist: $SRC_DIR" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
if [[ ! -d "$PATCHES_DIR" ]]; then
|
||||
echo "no patches dir at $PATCHES_DIR, nothing to apply"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
shopt -s nullglob
|
||||
patches=("$PATCHES_DIR"/*.patch)
|
||||
shopt -u nullglob
|
||||
|
||||
if [[ ${#patches[@]} -eq 0 ]]; then
|
||||
echo "no .patch files in $PATCHES_DIR, nothing to apply"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
cd "$SRC_DIR"
|
||||
|
||||
for patch in "${patches[@]}"; do
|
||||
echo "==> applying $patch"
|
||||
git apply --verbose "$patch"
|
||||
done
|
||||
|
||||
echo "all turboquant patches applied successfully"
|
||||
57
backend/cpp/turboquant/package.sh
Executable file
57
backend/cpp/turboquant/package.sh
Executable file
@@ -0,0 +1,57 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Script to copy the appropriate libraries based on architecture
|
||||
# This script is used in the final stage of the Dockerfile
|
||||
|
||||
set -e
|
||||
|
||||
CURDIR=$(dirname "$(realpath $0)")
|
||||
REPO_ROOT="${CURDIR}/../../.."
|
||||
|
||||
# Create lib directory
|
||||
mkdir -p $CURDIR/package/lib
|
||||
|
||||
cp -avrf $CURDIR/turboquant-* $CURDIR/package/
|
||||
cp -rfv $CURDIR/run.sh $CURDIR/package/
|
||||
|
||||
# Detect architecture and copy appropriate libraries
|
||||
if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then
|
||||
# x86_64 architecture
|
||||
echo "Detected x86_64 architecture, copying x86_64 libraries..."
|
||||
cp -arfLv /lib64/ld-linux-x86-64.so.2 $CURDIR/package/lib/ld.so
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libc.so.6 $CURDIR/package/lib/libc.so.6
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libgcc_s.so.1 $CURDIR/package/lib/libgcc_s.so.1
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libstdc++.so.6 $CURDIR/package/lib/libstdc++.so.6
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libm.so.6 $CURDIR/package/lib/libm.so.6
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libgomp.so.1 $CURDIR/package/lib/libgomp.so.1
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libdl.so.2 $CURDIR/package/lib/libdl.so.2
|
||||
cp -arfLv /lib/x86_64-linux-gnu/librt.so.1 $CURDIR/package/lib/librt.so.1
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libpthread.so.0 $CURDIR/package/lib/libpthread.so.0
|
||||
elif [ -f "/lib/ld-linux-aarch64.so.1" ]; then
|
||||
# ARM64 architecture
|
||||
echo "Detected ARM64 architecture, copying ARM64 libraries..."
|
||||
cp -arfLv /lib/ld-linux-aarch64.so.1 $CURDIR/package/lib/ld.so
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libc.so.6 $CURDIR/package/lib/libc.so.6
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libgcc_s.so.1 $CURDIR/package/lib/libgcc_s.so.1
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libstdc++.so.6 $CURDIR/package/lib/libstdc++.so.6
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libm.so.6 $CURDIR/package/lib/libm.so.6
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libgomp.so.1 $CURDIR/package/lib/libgomp.so.1
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libdl.so.2 $CURDIR/package/lib/libdl.so.2
|
||||
cp -arfLv /lib/aarch64-linux-gnu/librt.so.1 $CURDIR/package/lib/librt.so.1
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libpthread.so.0 $CURDIR/package/lib/libpthread.so.0
|
||||
else
|
||||
echo "Error: Could not detect architecture"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Package GPU libraries based on BUILD_TYPE
|
||||
GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh"
|
||||
if [ -f "$GPU_LIB_SCRIPT" ]; then
|
||||
echo "Packaging GPU libraries for BUILD_TYPE=${BUILD_TYPE:-cpu}..."
|
||||
source "$GPU_LIB_SCRIPT" "$CURDIR/package/lib"
|
||||
package_gpu_libs
|
||||
fi
|
||||
|
||||
echo "Packaging completed successfully"
|
||||
ls -liah $CURDIR/package/
|
||||
ls -liah $CURDIR/package/lib/
|
||||
80
backend/cpp/turboquant/patch-grpc-server.sh
Executable file
80
backend/cpp/turboquant/patch-grpc-server.sh
Executable file
@@ -0,0 +1,80 @@
|
||||
#!/bin/bash
|
||||
# Patch the shared backend/cpp/llama-cpp/grpc-server.cpp *copy* used by the
|
||||
# turboquant build to account for two gaps between upstream and the fork:
|
||||
#
|
||||
# 1. Augment the kv_cache_types[] allow-list so `LoadModel` accepts the
|
||||
# fork-specific `turbo2` / `turbo3` / `turbo4` cache types.
|
||||
# 2. Replace `get_media_marker()` (added upstream in ggml-org/llama.cpp#21962,
|
||||
# server-side random per-instance marker) with the legacy "<__media__>"
|
||||
# literal. The fork branched before that PR, so server-common.cpp has no
|
||||
# get_media_marker symbol. The fork's mtmd_default_marker() still returns
|
||||
# "<__media__>", and Go-side tooling falls back to that sentinel when the
|
||||
# backend does not expose media_marker, so substituting the literal keeps
|
||||
# behavior identical on the turboquant path.
|
||||
#
|
||||
# We patch the *copy* sitting in turboquant-<flavor>-build/, never the original
|
||||
# under backend/cpp/llama-cpp/, so the stock llama-cpp build keeps compiling
|
||||
# against vanilla upstream.
|
||||
#
|
||||
# Idempotent: skips each insertion if its marker is already present (so re-runs
|
||||
# of the same build dir don't double-insert).
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
if [[ $# -ne 1 ]]; then
|
||||
echo "usage: $0 <grpc-server.cpp>" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
SRC=$1
|
||||
|
||||
if [[ ! -f "$SRC" ]]; then
|
||||
echo "grpc-server.cpp not found at $SRC" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
if grep -q 'GGML_TYPE_TURBO2_0' "$SRC"; then
|
||||
echo "==> $SRC already has TurboQuant cache types, skipping KV allow-list patch"
|
||||
else
|
||||
echo "==> patching $SRC to allow turbo2/turbo3/turbo4 KV-cache types"
|
||||
|
||||
# Insert the three TURBO entries right after the first ` GGML_TYPE_Q5_1,`
|
||||
# line (the kv_cache_types[] allow-list). Using awk because the builder image
|
||||
# does not ship python3, and GNU sed's multi-line `a\` quoting is awkward.
|
||||
awk '
|
||||
/^ GGML_TYPE_Q5_1,$/ && !done {
|
||||
print
|
||||
print " // turboquant fork extras — added by patch-grpc-server.sh"
|
||||
print " GGML_TYPE_TURBO2_0,"
|
||||
print " GGML_TYPE_TURBO3_0,"
|
||||
print " GGML_TYPE_TURBO4_0,"
|
||||
done = 1
|
||||
next
|
||||
}
|
||||
{ print }
|
||||
END {
|
||||
if (!done) {
|
||||
print "patch-grpc-server.sh: anchor ` GGML_TYPE_Q5_1,` not found" > "/dev/stderr"
|
||||
exit 1
|
||||
}
|
||||
}
|
||||
' "$SRC" > "$SRC.tmp"
|
||||
mv "$SRC.tmp" "$SRC"
|
||||
|
||||
echo "==> KV allow-list patch OK"
|
||||
fi
|
||||
|
||||
if grep -q 'get_media_marker()' "$SRC"; then
|
||||
echo "==> patching $SRC to replace get_media_marker() with legacy \"<__media__>\" literal"
|
||||
# Only one call site today (ModelMetadata), but replace all occurrences to
|
||||
# stay robust if upstream adds more. Use a temp file to avoid relying on
|
||||
# sed -i portability (the builder image uses GNU sed, but keeping this
|
||||
# consistent with the awk block above).
|
||||
sed 's/get_media_marker()/"<__media__>"/g' "$SRC" > "$SRC.tmp"
|
||||
mv "$SRC.tmp" "$SRC"
|
||||
echo "==> get_media_marker() substitution OK"
|
||||
else
|
||||
echo "==> $SRC has no get_media_marker() call, skipping media-marker patch"
|
||||
fi
|
||||
|
||||
echo "==> all patches applied"
|
||||
65
backend/cpp/turboquant/run.sh
Executable file
65
backend/cpp/turboquant/run.sh
Executable file
@@ -0,0 +1,65 @@
|
||||
#!/bin/bash
|
||||
set -ex
|
||||
|
||||
# Get the absolute current dir where the script is located
|
||||
CURDIR=$(dirname "$(realpath $0)")
|
||||
|
||||
cd /
|
||||
|
||||
echo "CPU info:"
|
||||
grep -e "model\sname" /proc/cpuinfo | head -1
|
||||
grep -e "flags" /proc/cpuinfo | head -1
|
||||
|
||||
BINARY=turboquant-fallback
|
||||
|
||||
if grep -q -e "\savx\s" /proc/cpuinfo ; then
|
||||
echo "CPU: AVX found OK"
|
||||
if [ -e $CURDIR/turboquant-avx ]; then
|
||||
BINARY=turboquant-avx
|
||||
fi
|
||||
fi
|
||||
|
||||
if grep -q -e "\savx2\s" /proc/cpuinfo ; then
|
||||
echo "CPU: AVX2 found OK"
|
||||
if [ -e $CURDIR/turboquant-avx2 ]; then
|
||||
BINARY=turboquant-avx2
|
||||
fi
|
||||
fi
|
||||
|
||||
# Check avx 512
|
||||
if grep -q -e "\savx512f\s" /proc/cpuinfo ; then
|
||||
echo "CPU: AVX512F found OK"
|
||||
if [ -e $CURDIR/turboquant-avx512 ]; then
|
||||
BINARY=turboquant-avx512
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ -n "$LLAMACPP_GRPC_SERVERS" ]; then
|
||||
if [ -e $CURDIR/turboquant-grpc ]; then
|
||||
BINARY=turboquant-grpc
|
||||
fi
|
||||
fi
|
||||
|
||||
# Extend ld library path with the dir where this script is located/lib
|
||||
if [ "$(uname)" == "Darwin" ]; then
|
||||
export DYLD_LIBRARY_PATH=$CURDIR/lib:$DYLD_LIBRARY_PATH
|
||||
else
|
||||
export LD_LIBRARY_PATH=$CURDIR/lib:$LD_LIBRARY_PATH
|
||||
# Tell rocBLAS where to find TensileLibrary data (GPU kernel tuning files)
|
||||
if [ -d "$CURDIR/lib/rocblas/library" ]; then
|
||||
export ROCBLAS_TENSILE_LIBPATH=$CURDIR/lib/rocblas/library
|
||||
fi
|
||||
fi
|
||||
|
||||
# If there is a lib/ld.so, use it
|
||||
if [ -f $CURDIR/lib/ld.so ]; then
|
||||
echo "Using lib/ld.so"
|
||||
echo "Using binary: $BINARY"
|
||||
exec $CURDIR/lib/ld.so $CURDIR/$BINARY "$@"
|
||||
fi
|
||||
|
||||
echo "Using binary: $BINARY"
|
||||
exec $CURDIR/$BINARY "$@"
|
||||
|
||||
# We should never reach this point, however just in case we do, run fallback
|
||||
exec $CURDIR/turboquant-fallback "$@"
|
||||
@@ -4,7 +4,6 @@ package main
|
||||
// It is meant to be used by the main executable that is the server for the specific backend type (falcon, gpt3, etc)
|
||||
import (
|
||||
"container/heap"
|
||||
"errors"
|
||||
"fmt"
|
||||
"math"
|
||||
"slices"
|
||||
@@ -100,9 +99,16 @@ func sortIntoKeySlicese(keys []*pb.StoresKey) [][]float32 {
|
||||
}
|
||||
|
||||
func (s *Store) Load(opts *pb.ModelOptions) error {
|
||||
if opts.Model != "" {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
// local-store is an in-memory vector store with no on-disk artefact to
|
||||
// load — opts.Model is just a namespace identifier. The old `!= ""` guard
|
||||
// rejected any non-empty model name with "not implemented", which broke
|
||||
// callers that pass a namespace to isolate embedding spaces (face vs.
|
||||
// voice biometrics both go through local-store but need distinct stores
|
||||
// so ArcFace 512-D and ECAPA-TDNN 192-D don't collide). Namespace
|
||||
// isolation is already handled upstream: ModelLoader spawns a fresh
|
||||
// local-store process per (backend, model) tuple, so each namespace is
|
||||
// its own Store{} instance. Nothing to do here beyond accepting the load.
|
||||
_ = opts
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
11
backend/go/sherpa-onnx/.gitignore
vendored
Normal file
11
backend/go/sherpa-onnx/.gitignore
vendored
Normal file
@@ -0,0 +1,11 @@
|
||||
.cache/
|
||||
sources/
|
||||
build*/
|
||||
package/
|
||||
backend-assets/
|
||||
sherpa-onnx
|
||||
*.so
|
||||
compile_commands.json
|
||||
sherpa-onnx-whisper-*
|
||||
vits-ljs/
|
||||
streaming-zipformer-en/
|
||||
120
backend/go/sherpa-onnx/Makefile
Normal file
120
backend/go/sherpa-onnx/Makefile
Normal file
@@ -0,0 +1,120 @@
|
||||
CURRENT_DIR=$(abspath ./)
|
||||
GOCMD=go
|
||||
|
||||
ONNX_VERSION?=1.24.4
|
||||
# v1.12.39 — includes upstream's onnxruntime 1.24.4 bump (#3501). Earlier
|
||||
# pinned commits only support onnxruntime 1.23.2, which has no CUDA 13
|
||||
# pre-built tarball, blocking the -gpu-nvidia-cuda-13 build matrix entry.
|
||||
SHERPA_COMMIT?=7288d15e3e31a7bd589b2ba88828d521e7a6b140
|
||||
ONNX_ARCH?=x64
|
||||
ONNX_OS?=linux
|
||||
|
||||
ifneq (,$(findstring aarch64,$(shell uname -m)))
|
||||
ONNX_ARCH=aarch64
|
||||
endif
|
||||
|
||||
ifeq ($(OS),Darwin)
|
||||
ONNX_OS=osx
|
||||
ifneq (,$(findstring aarch64,$(shell uname -m)))
|
||||
ONNX_ARCH=arm64
|
||||
else ifneq (,$(findstring arm64,$(shell uname -m)))
|
||||
ONNX_ARCH=arm64
|
||||
else
|
||||
ONNX_ARCH=x86_64
|
||||
endif
|
||||
endif
|
||||
|
||||
# Upstream onnxruntime ships CUDA 12 and CUDA 13 variants under different
|
||||
# names: -gpu-<ver>.tgz for CUDA 12, -gpu_cuda13-<ver>.tgz for CUDA 13
|
||||
# (note underscore vs dash). CUDA 13 tarballs only exist from 1.24.x onward.
|
||||
ifeq ($(BUILD_TYPE),cublas)
|
||||
SHERPA_GPU=ON
|
||||
ONNX_PROVIDER=cuda
|
||||
ifeq ($(CUDA_MAJOR_VERSION),13)
|
||||
ONNX_VARIANT=-gpu_cuda13
|
||||
else
|
||||
ONNX_VARIANT=-gpu
|
||||
endif
|
||||
else
|
||||
ONNX_VARIANT=
|
||||
SHERPA_GPU=OFF
|
||||
ONNX_PROVIDER=cpu
|
||||
endif
|
||||
|
||||
JOBS?=$(shell nproc --ignore=1 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4)
|
||||
|
||||
sources/onnxruntime:
|
||||
mkdir -p sources/onnxruntime
|
||||
curl -L https://github.com/microsoft/onnxruntime/releases/download/v$(ONNX_VERSION)/onnxruntime-$(ONNX_OS)-$(ONNX_ARCH)$(ONNX_VARIANT)-$(ONNX_VERSION).tgz \
|
||||
-o sources/onnxruntime/onnxruntime.tgz
|
||||
cd sources/onnxruntime && tar -xf onnxruntime.tgz --strip-components=1 && rm onnxruntime.tgz
|
||||
|
||||
sources/sherpa-onnx: sources/onnxruntime
|
||||
git clone https://github.com/k2-fsa/sherpa-onnx.git sources/sherpa-onnx
|
||||
cd sources/sherpa-onnx && git checkout $(SHERPA_COMMIT)
|
||||
mkdir -p sources/sherpa-onnx/build
|
||||
# sherpa-onnx's cmake detects a pre-installed onnxruntime via the
|
||||
# SHERPA_ONNXRUNTIME_{INCLUDE,LIB}_DIR env vars (not via -D flags).
|
||||
# Point them at our locally-downloaded Microsoft tarball — without
|
||||
# this, sherpa-onnx falls through to download_onnxruntime() which
|
||||
# fetches from csukuangfj/onnxruntime-libs. For the GPU 1.24.4
|
||||
# build that release mirror publishes `-patched.zip` instead of the
|
||||
# expected `.tgz`, so the download 404s and the build fails.
|
||||
cd sources/sherpa-onnx/build && \
|
||||
SHERPA_ONNXRUNTIME_INCLUDE_DIR=$(CURRENT_DIR)/sources/onnxruntime/include \
|
||||
SHERPA_ONNXRUNTIME_LIB_DIR=$(CURRENT_DIR)/sources/onnxruntime/lib \
|
||||
cmake \
|
||||
-DCMAKE_BUILD_TYPE=Release \
|
||||
-DCMAKE_C_FLAGS="-Wno-error=format-security" \
|
||||
-DCMAKE_CXX_FLAGS="-Wno-error=format-security" \
|
||||
-DSHERPA_ONNX_ENABLE_GPU=$(SHERPA_GPU) \
|
||||
-DSHERPA_ONNX_ENABLE_TTS=ON \
|
||||
-DSHERPA_ONNX_ENABLE_BINARY=OFF \
|
||||
-DSHERPA_ONNX_ENABLE_PYTHON=OFF \
|
||||
-DSHERPA_ONNX_ENABLE_TESTS=OFF \
|
||||
-DSHERPA_ONNX_ENABLE_C_API=ON \
|
||||
-DBUILD_SHARED_LIBS=ON \
|
||||
-DSHERPA_ONNX_USE_PRE_INSTALLED_ONNXRUNTIME_IF_AVAILABLE=ON \
|
||||
..
|
||||
cd sources/sherpa-onnx/build && make -j$(JOBS)
|
||||
|
||||
backend-assets/lib: sources/sherpa-onnx sources/onnxruntime
|
||||
mkdir -p backend-assets/lib
|
||||
cp -rfLv sources/onnxruntime/lib/* backend-assets/lib/
|
||||
cp -rfLv sources/sherpa-onnx/build/lib/*.so* backend-assets/lib/ 2>/dev/null || true
|
||||
cp -rfLv sources/sherpa-onnx/build/lib/*.dylib backend-assets/lib/ 2>/dev/null || true
|
||||
|
||||
# libsherpa-shim wraps sherpa-onnx's nested config structs and TTS
|
||||
# callback plumbing behind a purego-friendly API: opaque handles plus
|
||||
# fixed-signature setters/getters/trampoline. Plain C compile — no cgo.
|
||||
SHIM_EXT=so
|
||||
ifeq ($(OS),Darwin)
|
||||
SHIM_EXT=dylib
|
||||
endif
|
||||
|
||||
backend-assets/lib/libsherpa-shim.$(SHIM_EXT): csrc/shim.c csrc/shim.h backend-assets/lib
|
||||
$(CC) -shared -fPIC -O2 \
|
||||
-I$(CURRENT_DIR)/sources/sherpa-onnx/sherpa-onnx/c-api \
|
||||
-o $@ csrc/shim.c \
|
||||
-L$(CURRENT_DIR)/backend-assets/lib \
|
||||
-lsherpa-onnx-c-api \
|
||||
-Wl,-rpath,'$$ORIGIN'
|
||||
|
||||
sherpa-onnx: backend-assets/lib backend-assets/lib/libsherpa-shim.$(SHIM_EXT)
|
||||
CGO_ENABLED=0 $(GOCMD) build \
|
||||
-ldflags "$(LD_FLAGS) -X main.onnxProvider=$(ONNX_PROVIDER)" \
|
||||
-tags "$(GO_TAGS)" -o sherpa-onnx ./
|
||||
|
||||
package:
|
||||
bash package.sh
|
||||
|
||||
build: sherpa-onnx package
|
||||
|
||||
clean:
|
||||
rm -rf sherpa-onnx sources/ backend-assets/ package/ vits-ljs/ sherpa-onnx-whisper-*/
|
||||
|
||||
test: sherpa-onnx
|
||||
LD_LIBRARY_PATH=$(CURRENT_DIR)/backend-assets/lib \
|
||||
bash test.sh
|
||||
|
||||
.PHONY: build package clean test
|
||||
1249
backend/go/sherpa-onnx/backend.go
Normal file
1249
backend/go/sherpa-onnx/backend.go
Normal file
File diff suppressed because it is too large
Load Diff
169
backend/go/sherpa-onnx/backend_test.go
Normal file
169
backend/go/sherpa-onnx/backend_test.go
Normal file
@@ -0,0 +1,169 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
func TestSherpaBackend(t *testing.T) {
|
||||
RegisterFailHandler(Fail)
|
||||
RunSpecs(t, "Sherpa-ONNX Backend Suite")
|
||||
}
|
||||
|
||||
// Load libsherpa-shim + libsherpa-onnx-c-api via purego before any spec
|
||||
// runs — otherwise any Load/TTS/VAD/AudioTranscription call hits a nil
|
||||
// function pointer. LD_LIBRARY_PATH must contain the directory holding
|
||||
// both .so files; test.sh sets this.
|
||||
var _ = BeforeSuite(func() {
|
||||
Expect(loadSherpaLibs()).To(Succeed())
|
||||
})
|
||||
|
||||
var _ = Describe("Sherpa-ONNX", func() {
|
||||
Context("lifecycle", func() {
|
||||
It("is locking (C API is not thread safe)", func() {
|
||||
Expect((&SherpaBackend{}).Locking()).To(BeTrue())
|
||||
})
|
||||
|
||||
It("errors loading a non-existent model", func() {
|
||||
tmpDir, err := os.MkdirTemp("", "sherpa-test-nonexistent")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
err = (&SherpaBackend{}).Load(&pb.ModelOptions{
|
||||
ModelFile: filepath.Join(tmpDir, "non-existent-model.onnx"),
|
||||
})
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("errors loading a non-existent ASR model", func() {
|
||||
tmpDir, err := os.MkdirTemp("", "sherpa-test-asr")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
err = (&SherpaBackend{}).Load(&pb.ModelOptions{
|
||||
ModelFile: filepath.Join(tmpDir, "model.onnx"),
|
||||
Type: "asr",
|
||||
})
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("dispatches Load by Type", func() {
|
||||
tmpDir, err := os.MkdirTemp("", "sherpa-test-dispatch")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
modelFile := filepath.Join(tmpDir, "model.onnx")
|
||||
for _, typ := range []string{"", "asr", "vad"} {
|
||||
err := (&SherpaBackend{}).Load(&pb.ModelOptions{ModelFile: modelFile, Type: typ})
|
||||
Expect(err).To(HaveOccurred(), "Type=%q", typ)
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
Context("method errors without loaded model", func() {
|
||||
It("rejects TTS", func() {
|
||||
tmpDir, err := os.MkdirTemp("", "sherpa-test-tts")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
err = (&SherpaBackend{}).TTS(&pb.TTSRequest{
|
||||
Text: "should fail — no model loaded",
|
||||
Dst: filepath.Join(tmpDir, "output.wav"),
|
||||
})
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("rejects AudioTranscription", func() {
|
||||
_, err := (&SherpaBackend{}).AudioTranscription(&pb.TranscriptRequest{
|
||||
Dst: "/tmp/nonexistent.wav",
|
||||
})
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("rejects VAD", func() {
|
||||
_, err := (&SherpaBackend{}).VAD(&pb.VADRequest{
|
||||
Audio: []float32{0.1, 0.2, 0.3},
|
||||
})
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
})
|
||||
|
||||
Context("type detection", func() {
|
||||
DescribeTable("isASRType",
|
||||
func(input string, want bool) {
|
||||
Expect(isASRType(input)).To(Equal(want))
|
||||
},
|
||||
Entry("asr", "asr", true),
|
||||
Entry("ASR", "ASR", true),
|
||||
Entry("Asr", "Asr", true),
|
||||
Entry("transcription", "transcription", true),
|
||||
Entry("Transcription", "Transcription", true),
|
||||
Entry("transcribe", "transcribe", true),
|
||||
Entry("Transcribe", "Transcribe", true),
|
||||
Entry("tts", "tts", false),
|
||||
Entry("empty", "", false),
|
||||
Entry("other", "other", false),
|
||||
Entry("vad", "vad", false),
|
||||
)
|
||||
|
||||
DescribeTable("isVADType",
|
||||
func(input string, want bool) {
|
||||
Expect(isVADType(input)).To(Equal(want))
|
||||
},
|
||||
Entry("vad", "vad", true),
|
||||
Entry("VAD", "VAD", true),
|
||||
Entry("Vad", "Vad", true),
|
||||
Entry("asr", "asr", false),
|
||||
Entry("tts", "tts", false),
|
||||
Entry("empty", "", false),
|
||||
Entry("other", "other", false),
|
||||
)
|
||||
})
|
||||
|
||||
Context("option parsing", func() {
|
||||
It("parses float options with fallback on bad input", func() {
|
||||
opts := &pb.ModelOptions{Options: []string{
|
||||
"vad.threshold=0.3",
|
||||
"tts.length_scale=1.25",
|
||||
"bad.number=not-a-float",
|
||||
}}
|
||||
Expect(findOptionFloat(opts, "vad.threshold=", 0.5)).To(BeNumerically("~", 0.3, 1e-6))
|
||||
Expect(findOptionFloat(opts, "tts.length_scale=", 1.0)).To(BeNumerically("~", 1.25, 1e-6))
|
||||
Expect(findOptionFloat(opts, "missing.key=", 0.7)).To(BeNumerically("~", 0.7, 1e-6))
|
||||
Expect(findOptionFloat(opts, "bad.number=", 9.9)).To(BeNumerically("~", 9.9, 1e-6))
|
||||
})
|
||||
|
||||
It("parses int options with fallback on bad input", func() {
|
||||
opts := &pb.ModelOptions{Options: []string{
|
||||
"asr.sample_rate=22050",
|
||||
"online.chunk_samples=800",
|
||||
"bad.int=4.2",
|
||||
}}
|
||||
Expect(findOptionInt(opts, "asr.sample_rate=", 16000)).To(Equal(int32(22050)))
|
||||
Expect(findOptionInt(opts, "online.chunk_samples=", 1600)).To(Equal(int32(800)))
|
||||
Expect(findOptionInt(opts, "missing.key=", 42)).To(Equal(int32(42)))
|
||||
Expect(findOptionInt(opts, "bad.int=", 100)).To(Equal(int32(100)))
|
||||
})
|
||||
|
||||
It("parses bool options (0/1, true/false, yes/no, on/off)", func() {
|
||||
opts := &pb.ModelOptions{Options: []string{
|
||||
"online.enable_endpoint=0",
|
||||
"asr.sense_voice.use_itn=True",
|
||||
"feature.on=yes",
|
||||
"feature.off=Off",
|
||||
"feature.bad=maybe",
|
||||
}}
|
||||
Expect(findOptionBool(opts, "online.enable_endpoint=", 1)).To(Equal(int32(0)))
|
||||
Expect(findOptionBool(opts, "asr.sense_voice.use_itn=", 0)).To(Equal(int32(1)))
|
||||
Expect(findOptionBool(opts, "feature.on=", 0)).To(Equal(int32(1)))
|
||||
Expect(findOptionBool(opts, "feature.off=", 1)).To(Equal(int32(0)))
|
||||
Expect(findOptionBool(opts, "feature.bad=", 1)).To(Equal(int32(1)))
|
||||
Expect(findOptionBool(opts, "missing.key=", 1)).To(Equal(int32(1)))
|
||||
})
|
||||
})
|
||||
})
|
||||
325
backend/go/sherpa-onnx/csrc/shim.c
Normal file
325
backend/go/sherpa-onnx/csrc/shim.c
Normal file
@@ -0,0 +1,325 @@
|
||||
#include "shim.h"
|
||||
#include "c-api.h"
|
||||
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
|
||||
// Replace the char* field pointed to by `slot` with a strdup of `s`
|
||||
// (or NULL if s is NULL). Frees any prior value. Silently no-ops when
|
||||
// strdup fails — the caller will see a Create* failure downstream.
|
||||
static void shim_set_str(const char **slot, const char *s) {
|
||||
free((char *)*slot);
|
||||
*slot = s ? strdup(s) : NULL;
|
||||
}
|
||||
|
||||
// ==================================================================
|
||||
// VAD config
|
||||
// ==================================================================
|
||||
|
||||
void *sherpa_shim_vad_config_new(void) {
|
||||
return calloc(1, sizeof(SherpaOnnxVadModelConfig));
|
||||
}
|
||||
|
||||
void sherpa_shim_vad_config_free(void *h) {
|
||||
if (!h) return;
|
||||
SherpaOnnxVadModelConfig *c = (SherpaOnnxVadModelConfig *)h;
|
||||
free((char *)c->silero_vad.model);
|
||||
free((char *)c->provider);
|
||||
free(c);
|
||||
}
|
||||
|
||||
void sherpa_shim_vad_config_set_silero_model(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxVadModelConfig *)h)->silero_vad.model, v);
|
||||
}
|
||||
void sherpa_shim_vad_config_set_silero_threshold(void *h, float v) {
|
||||
((SherpaOnnxVadModelConfig *)h)->silero_vad.threshold = v;
|
||||
}
|
||||
void sherpa_shim_vad_config_set_silero_min_silence_duration(void *h, float v) {
|
||||
((SherpaOnnxVadModelConfig *)h)->silero_vad.min_silence_duration = v;
|
||||
}
|
||||
void sherpa_shim_vad_config_set_silero_min_speech_duration(void *h, float v) {
|
||||
((SherpaOnnxVadModelConfig *)h)->silero_vad.min_speech_duration = v;
|
||||
}
|
||||
void sherpa_shim_vad_config_set_silero_window_size(void *h, int32_t v) {
|
||||
((SherpaOnnxVadModelConfig *)h)->silero_vad.window_size = v;
|
||||
}
|
||||
void sherpa_shim_vad_config_set_silero_max_speech_duration(void *h, float v) {
|
||||
((SherpaOnnxVadModelConfig *)h)->silero_vad.max_speech_duration = v;
|
||||
}
|
||||
void sherpa_shim_vad_config_set_sample_rate(void *h, int32_t v) {
|
||||
((SherpaOnnxVadModelConfig *)h)->sample_rate = v;
|
||||
}
|
||||
void sherpa_shim_vad_config_set_num_threads(void *h, int32_t v) {
|
||||
((SherpaOnnxVadModelConfig *)h)->num_threads = v;
|
||||
}
|
||||
void sherpa_shim_vad_config_set_provider(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxVadModelConfig *)h)->provider, v);
|
||||
}
|
||||
void sherpa_shim_vad_config_set_debug(void *h, int32_t v) {
|
||||
((SherpaOnnxVadModelConfig *)h)->debug = v;
|
||||
}
|
||||
|
||||
void *sherpa_shim_create_vad(void *h, float buffer_size_seconds) {
|
||||
return (void *)SherpaOnnxCreateVoiceActivityDetector(
|
||||
(const SherpaOnnxVadModelConfig *)h, buffer_size_seconds);
|
||||
}
|
||||
|
||||
// ==================================================================
|
||||
// Offline TTS config (VITS)
|
||||
// ==================================================================
|
||||
|
||||
void *sherpa_shim_tts_config_new(void) {
|
||||
return calloc(1, sizeof(SherpaOnnxOfflineTtsConfig));
|
||||
}
|
||||
|
||||
void sherpa_shim_tts_config_free(void *h) {
|
||||
if (!h) return;
|
||||
SherpaOnnxOfflineTtsConfig *c = (SherpaOnnxOfflineTtsConfig *)h;
|
||||
free((char *)c->model.vits.model);
|
||||
free((char *)c->model.vits.tokens);
|
||||
free((char *)c->model.vits.lexicon);
|
||||
free((char *)c->model.vits.data_dir);
|
||||
free((char *)c->model.provider);
|
||||
free(c);
|
||||
}
|
||||
|
||||
void sherpa_shim_tts_config_set_vits_model(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOfflineTtsConfig *)h)->model.vits.model, v);
|
||||
}
|
||||
void sherpa_shim_tts_config_set_vits_tokens(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOfflineTtsConfig *)h)->model.vits.tokens, v);
|
||||
}
|
||||
void sherpa_shim_tts_config_set_vits_lexicon(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOfflineTtsConfig *)h)->model.vits.lexicon, v);
|
||||
}
|
||||
void sherpa_shim_tts_config_set_vits_data_dir(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOfflineTtsConfig *)h)->model.vits.data_dir, v);
|
||||
}
|
||||
void sherpa_shim_tts_config_set_vits_noise_scale(void *h, float v) {
|
||||
((SherpaOnnxOfflineTtsConfig *)h)->model.vits.noise_scale = v;
|
||||
}
|
||||
void sherpa_shim_tts_config_set_vits_noise_scale_w(void *h, float v) {
|
||||
((SherpaOnnxOfflineTtsConfig *)h)->model.vits.noise_scale_w = v;
|
||||
}
|
||||
void sherpa_shim_tts_config_set_vits_length_scale(void *h, float v) {
|
||||
((SherpaOnnxOfflineTtsConfig *)h)->model.vits.length_scale = v;
|
||||
}
|
||||
void sherpa_shim_tts_config_set_num_threads(void *h, int32_t v) {
|
||||
((SherpaOnnxOfflineTtsConfig *)h)->model.num_threads = v;
|
||||
}
|
||||
void sherpa_shim_tts_config_set_debug(void *h, int32_t v) {
|
||||
((SherpaOnnxOfflineTtsConfig *)h)->model.debug = v;
|
||||
}
|
||||
void sherpa_shim_tts_config_set_provider(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOfflineTtsConfig *)h)->model.provider, v);
|
||||
}
|
||||
void sherpa_shim_tts_config_set_max_num_sentences(void *h, int32_t v) {
|
||||
((SherpaOnnxOfflineTtsConfig *)h)->max_num_sentences = v;
|
||||
}
|
||||
|
||||
void *sherpa_shim_create_offline_tts(void *h) {
|
||||
return (void *)SherpaOnnxCreateOfflineTts(
|
||||
(const SherpaOnnxOfflineTtsConfig *)h);
|
||||
}
|
||||
|
||||
// ==================================================================
|
||||
// Offline recognizer config
|
||||
// ==================================================================
|
||||
|
||||
void *sherpa_shim_offline_recog_config_new(void) {
|
||||
return calloc(1, sizeof(SherpaOnnxOfflineRecognizerConfig));
|
||||
}
|
||||
|
||||
void sherpa_shim_offline_recog_config_free(void *h) {
|
||||
if (!h) return;
|
||||
SherpaOnnxOfflineRecognizerConfig *c = (SherpaOnnxOfflineRecognizerConfig *)h;
|
||||
free((char *)c->model_config.provider);
|
||||
free((char *)c->model_config.tokens);
|
||||
free((char *)c->model_config.whisper.encoder);
|
||||
free((char *)c->model_config.whisper.decoder);
|
||||
free((char *)c->model_config.whisper.language);
|
||||
free((char *)c->model_config.whisper.task);
|
||||
free((char *)c->model_config.paraformer.model);
|
||||
free((char *)c->model_config.sense_voice.model);
|
||||
free((char *)c->model_config.sense_voice.language);
|
||||
free((char *)c->model_config.omnilingual.model);
|
||||
free((char *)c->decoding_method);
|
||||
free(c);
|
||||
}
|
||||
|
||||
void sherpa_shim_offline_recog_config_set_num_threads(void *h, int32_t v) {
|
||||
((SherpaOnnxOfflineRecognizerConfig *)h)->model_config.num_threads = v;
|
||||
}
|
||||
void sherpa_shim_offline_recog_config_set_debug(void *h, int32_t v) {
|
||||
((SherpaOnnxOfflineRecognizerConfig *)h)->model_config.debug = v;
|
||||
}
|
||||
void sherpa_shim_offline_recog_config_set_provider(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOfflineRecognizerConfig *)h)->model_config.provider, v);
|
||||
}
|
||||
void sherpa_shim_offline_recog_config_set_tokens(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOfflineRecognizerConfig *)h)->model_config.tokens, v);
|
||||
}
|
||||
void sherpa_shim_offline_recog_config_set_feat_sample_rate(void *h, int32_t v) {
|
||||
((SherpaOnnxOfflineRecognizerConfig *)h)->feat_config.sample_rate = v;
|
||||
}
|
||||
void sherpa_shim_offline_recog_config_set_feat_feature_dim(void *h, int32_t v) {
|
||||
((SherpaOnnxOfflineRecognizerConfig *)h)->feat_config.feature_dim = v;
|
||||
}
|
||||
void sherpa_shim_offline_recog_config_set_decoding_method(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOfflineRecognizerConfig *)h)->decoding_method, v);
|
||||
}
|
||||
void sherpa_shim_offline_recog_config_set_whisper_encoder(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOfflineRecognizerConfig *)h)->model_config.whisper.encoder, v);
|
||||
}
|
||||
void sherpa_shim_offline_recog_config_set_whisper_decoder(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOfflineRecognizerConfig *)h)->model_config.whisper.decoder, v);
|
||||
}
|
||||
void sherpa_shim_offline_recog_config_set_whisper_language(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOfflineRecognizerConfig *)h)->model_config.whisper.language, v);
|
||||
}
|
||||
void sherpa_shim_offline_recog_config_set_whisper_task(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOfflineRecognizerConfig *)h)->model_config.whisper.task, v);
|
||||
}
|
||||
void sherpa_shim_offline_recog_config_set_whisper_tail_paddings(void *h, int32_t v) {
|
||||
((SherpaOnnxOfflineRecognizerConfig *)h)->model_config.whisper.tail_paddings = v;
|
||||
}
|
||||
void sherpa_shim_offline_recog_config_set_paraformer_model(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOfflineRecognizerConfig *)h)->model_config.paraformer.model, v);
|
||||
}
|
||||
void sherpa_shim_offline_recog_config_set_sense_voice_model(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOfflineRecognizerConfig *)h)->model_config.sense_voice.model, v);
|
||||
}
|
||||
void sherpa_shim_offline_recog_config_set_sense_voice_language(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOfflineRecognizerConfig *)h)->model_config.sense_voice.language, v);
|
||||
}
|
||||
void sherpa_shim_offline_recog_config_set_sense_voice_use_itn(void *h, int32_t v) {
|
||||
((SherpaOnnxOfflineRecognizerConfig *)h)->model_config.sense_voice.use_itn = v;
|
||||
}
|
||||
void sherpa_shim_offline_recog_config_set_omnilingual_model(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOfflineRecognizerConfig *)h)->model_config.omnilingual.model, v);
|
||||
}
|
||||
|
||||
void *sherpa_shim_create_offline_recognizer(void *h) {
|
||||
return (void *)SherpaOnnxCreateOfflineRecognizer(
|
||||
(const SherpaOnnxOfflineRecognizerConfig *)h);
|
||||
}
|
||||
|
||||
// ==================================================================
|
||||
// Online recognizer config
|
||||
// ==================================================================
|
||||
|
||||
void *sherpa_shim_online_recog_config_new(void) {
|
||||
return calloc(1, sizeof(SherpaOnnxOnlineRecognizerConfig));
|
||||
}
|
||||
|
||||
void sherpa_shim_online_recog_config_free(void *h) {
|
||||
if (!h) return;
|
||||
SherpaOnnxOnlineRecognizerConfig *c = (SherpaOnnxOnlineRecognizerConfig *)h;
|
||||
free((char *)c->model_config.transducer.encoder);
|
||||
free((char *)c->model_config.transducer.decoder);
|
||||
free((char *)c->model_config.transducer.joiner);
|
||||
free((char *)c->model_config.tokens);
|
||||
free((char *)c->model_config.provider);
|
||||
free((char *)c->decoding_method);
|
||||
free(c);
|
||||
}
|
||||
|
||||
void sherpa_shim_online_recog_config_set_transducer_encoder(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOnlineRecognizerConfig *)h)->model_config.transducer.encoder, v);
|
||||
}
|
||||
void sherpa_shim_online_recog_config_set_transducer_decoder(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOnlineRecognizerConfig *)h)->model_config.transducer.decoder, v);
|
||||
}
|
||||
void sherpa_shim_online_recog_config_set_transducer_joiner(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOnlineRecognizerConfig *)h)->model_config.transducer.joiner, v);
|
||||
}
|
||||
void sherpa_shim_online_recog_config_set_tokens(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOnlineRecognizerConfig *)h)->model_config.tokens, v);
|
||||
}
|
||||
void sherpa_shim_online_recog_config_set_num_threads(void *h, int32_t v) {
|
||||
((SherpaOnnxOnlineRecognizerConfig *)h)->model_config.num_threads = v;
|
||||
}
|
||||
void sherpa_shim_online_recog_config_set_debug(void *h, int32_t v) {
|
||||
((SherpaOnnxOnlineRecognizerConfig *)h)->model_config.debug = v;
|
||||
}
|
||||
void sherpa_shim_online_recog_config_set_provider(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOnlineRecognizerConfig *)h)->model_config.provider, v);
|
||||
}
|
||||
void sherpa_shim_online_recog_config_set_feat_sample_rate(void *h, int32_t v) {
|
||||
((SherpaOnnxOnlineRecognizerConfig *)h)->feat_config.sample_rate = v;
|
||||
}
|
||||
void sherpa_shim_online_recog_config_set_feat_feature_dim(void *h, int32_t v) {
|
||||
((SherpaOnnxOnlineRecognizerConfig *)h)->feat_config.feature_dim = v;
|
||||
}
|
||||
void sherpa_shim_online_recog_config_set_decoding_method(void *h, const char *v) {
|
||||
shim_set_str(&((SherpaOnnxOnlineRecognizerConfig *)h)->decoding_method, v);
|
||||
}
|
||||
void sherpa_shim_online_recog_config_set_enable_endpoint(void *h, int32_t v) {
|
||||
((SherpaOnnxOnlineRecognizerConfig *)h)->enable_endpoint = v;
|
||||
}
|
||||
void sherpa_shim_online_recog_config_set_rule1_min_trailing_silence(void *h, float v) {
|
||||
((SherpaOnnxOnlineRecognizerConfig *)h)->rule1_min_trailing_silence = v;
|
||||
}
|
||||
void sherpa_shim_online_recog_config_set_rule2_min_trailing_silence(void *h, float v) {
|
||||
((SherpaOnnxOnlineRecognizerConfig *)h)->rule2_min_trailing_silence = v;
|
||||
}
|
||||
void sherpa_shim_online_recog_config_set_rule3_min_utterance_length(void *h, float v) {
|
||||
((SherpaOnnxOnlineRecognizerConfig *)h)->rule3_min_utterance_length = v;
|
||||
}
|
||||
|
||||
void *sherpa_shim_create_online_recognizer(void *h) {
|
||||
return (void *)SherpaOnnxCreateOnlineRecognizer(
|
||||
(const SherpaOnnxOnlineRecognizerConfig *)h);
|
||||
}
|
||||
|
||||
// ==================================================================
|
||||
// Result-struct accessors
|
||||
// ==================================================================
|
||||
|
||||
int32_t sherpa_shim_wave_sample_rate(const void *h) {
|
||||
return ((const SherpaOnnxWave *)h)->sample_rate;
|
||||
}
|
||||
int32_t sherpa_shim_wave_num_samples(const void *h) {
|
||||
return ((const SherpaOnnxWave *)h)->num_samples;
|
||||
}
|
||||
const float *sherpa_shim_wave_samples(const void *h) {
|
||||
return ((const SherpaOnnxWave *)h)->samples;
|
||||
}
|
||||
|
||||
const char *sherpa_shim_offline_result_text(const void *h) {
|
||||
return ((const SherpaOnnxOfflineRecognizerResult *)h)->text;
|
||||
}
|
||||
const char *sherpa_shim_online_result_text(const void *h) {
|
||||
return ((const SherpaOnnxOnlineRecognizerResult *)h)->text;
|
||||
}
|
||||
|
||||
int32_t sherpa_shim_generated_audio_sample_rate(const void *h) {
|
||||
return ((const SherpaOnnxGeneratedAudio *)h)->sample_rate;
|
||||
}
|
||||
int32_t sherpa_shim_generated_audio_n(const void *h) {
|
||||
return ((const SherpaOnnxGeneratedAudio *)h)->n;
|
||||
}
|
||||
const float *sherpa_shim_generated_audio_samples(const void *h) {
|
||||
return ((const SherpaOnnxGeneratedAudio *)h)->samples;
|
||||
}
|
||||
|
||||
int32_t sherpa_shim_speech_segment_start(const void *h) {
|
||||
return ((const SherpaOnnxSpeechSegment *)h)->start;
|
||||
}
|
||||
int32_t sherpa_shim_speech_segment_n(const void *h) {
|
||||
return ((const SherpaOnnxSpeechSegment *)h)->n;
|
||||
}
|
||||
|
||||
// ==================================================================
|
||||
// TTS streaming callback trampoline
|
||||
// ==================================================================
|
||||
|
||||
void *sherpa_shim_tts_generate_with_callback(
|
||||
void *tts, const char *text, int32_t sid, float speed,
|
||||
uintptr_t callback_ptr, uintptr_t user_data) {
|
||||
SherpaOnnxGeneratedAudioCallbackWithArg cb =
|
||||
(SherpaOnnxGeneratedAudioCallbackWithArg)callback_ptr;
|
||||
return (void *)SherpaOnnxOfflineTtsGenerateWithCallbackWithArg(
|
||||
(const SherpaOnnxOfflineTts *)tts, text, sid, speed, cb,
|
||||
(void *)user_data);
|
||||
}
|
||||
129
backend/go/sherpa-onnx/csrc/shim.h
Normal file
129
backend/go/sherpa-onnx/csrc/shim.h
Normal file
@@ -0,0 +1,129 @@
|
||||
#ifndef LOCALAI_SHERPA_ONNX_SHIM_H
|
||||
#define LOCALAI_SHERPA_ONNX_SHIM_H
|
||||
|
||||
#include <stdint.h>
|
||||
|
||||
// libsherpa-shim: purego-friendly wrapper around sherpa-onnx's C API.
|
||||
// Purego can't access C struct fields and can't route C callbacks to Go
|
||||
// funcs directly. Every function here is a fixed-signature trampoline
|
||||
// that replaces one field read/write or callback handoff that the Go
|
||||
// backend would otherwise have to do through cgo.
|
||||
//
|
||||
// String lifetime: setters strdup; _free walks every owned string and
|
||||
// frees it. Callers may discard their input buffers the moment a setter
|
||||
// returns.
|
||||
//
|
||||
// Opaque handles are `void *` in both directions. Nothing here holds a
|
||||
// reference across calls except config handles (freed via _free) and
|
||||
// sherpa-allocated results (freed via sherpa's own Destroy* entry
|
||||
// points, which Go calls through purego pass-through).
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
// --- VAD config -----------------------------------------------------
|
||||
void *sherpa_shim_vad_config_new(void);
|
||||
void sherpa_shim_vad_config_free(void *cfg);
|
||||
void sherpa_shim_vad_config_set_silero_model(void *cfg, const char *path);
|
||||
void sherpa_shim_vad_config_set_silero_threshold(void *cfg, float v);
|
||||
void sherpa_shim_vad_config_set_silero_min_silence_duration(void *cfg, float v);
|
||||
void sherpa_shim_vad_config_set_silero_min_speech_duration(void *cfg, float v);
|
||||
void sherpa_shim_vad_config_set_silero_window_size(void *cfg, int32_t v);
|
||||
void sherpa_shim_vad_config_set_silero_max_speech_duration(void *cfg, float v);
|
||||
void sherpa_shim_vad_config_set_sample_rate(void *cfg, int32_t v);
|
||||
void sherpa_shim_vad_config_set_num_threads(void *cfg, int32_t v);
|
||||
void sherpa_shim_vad_config_set_provider(void *cfg, const char *v);
|
||||
void sherpa_shim_vad_config_set_debug(void *cfg, int32_t v);
|
||||
void *sherpa_shim_create_vad(void *cfg, float buffer_size_seconds);
|
||||
|
||||
// --- Offline TTS config (VITS path — the only TTS family the backend uses) ---
|
||||
void *sherpa_shim_tts_config_new(void);
|
||||
void sherpa_shim_tts_config_free(void *cfg);
|
||||
void sherpa_shim_tts_config_set_vits_model(void *cfg, const char *v);
|
||||
void sherpa_shim_tts_config_set_vits_tokens(void *cfg, const char *v);
|
||||
void sherpa_shim_tts_config_set_vits_lexicon(void *cfg, const char *v);
|
||||
void sherpa_shim_tts_config_set_vits_data_dir(void *cfg, const char *v);
|
||||
void sherpa_shim_tts_config_set_vits_noise_scale(void *cfg, float v);
|
||||
void sherpa_shim_tts_config_set_vits_noise_scale_w(void *cfg, float v);
|
||||
void sherpa_shim_tts_config_set_vits_length_scale(void *cfg, float v);
|
||||
void sherpa_shim_tts_config_set_num_threads(void *cfg, int32_t v);
|
||||
void sherpa_shim_tts_config_set_debug(void *cfg, int32_t v);
|
||||
void sherpa_shim_tts_config_set_provider(void *cfg, const char *v);
|
||||
void sherpa_shim_tts_config_set_max_num_sentences(void *cfg, int32_t v);
|
||||
void *sherpa_shim_create_offline_tts(void *cfg);
|
||||
|
||||
// --- Offline recognizer config (Whisper / Paraformer / SenseVoice / Omnilingual) ---
|
||||
void *sherpa_shim_offline_recog_config_new(void);
|
||||
void sherpa_shim_offline_recog_config_free(void *cfg);
|
||||
void sherpa_shim_offline_recog_config_set_num_threads(void *cfg, int32_t v);
|
||||
void sherpa_shim_offline_recog_config_set_debug(void *cfg, int32_t v);
|
||||
void sherpa_shim_offline_recog_config_set_provider(void *cfg, const char *v);
|
||||
void sherpa_shim_offline_recog_config_set_tokens(void *cfg, const char *v);
|
||||
void sherpa_shim_offline_recog_config_set_feat_sample_rate(void *cfg, int32_t v);
|
||||
void sherpa_shim_offline_recog_config_set_feat_feature_dim(void *cfg, int32_t v);
|
||||
void sherpa_shim_offline_recog_config_set_decoding_method(void *cfg, const char *v);
|
||||
void sherpa_shim_offline_recog_config_set_whisper_encoder(void *cfg, const char *v);
|
||||
void sherpa_shim_offline_recog_config_set_whisper_decoder(void *cfg, const char *v);
|
||||
void sherpa_shim_offline_recog_config_set_whisper_language(void *cfg, const char *v);
|
||||
void sherpa_shim_offline_recog_config_set_whisper_task(void *cfg, const char *v);
|
||||
void sherpa_shim_offline_recog_config_set_whisper_tail_paddings(void *cfg, int32_t v);
|
||||
void sherpa_shim_offline_recog_config_set_paraformer_model(void *cfg, const char *v);
|
||||
void sherpa_shim_offline_recog_config_set_sense_voice_model(void *cfg, const char *v);
|
||||
void sherpa_shim_offline_recog_config_set_sense_voice_language(void *cfg, const char *v);
|
||||
void sherpa_shim_offline_recog_config_set_sense_voice_use_itn(void *cfg, int32_t v);
|
||||
void sherpa_shim_offline_recog_config_set_omnilingual_model(void *cfg, const char *v);
|
||||
void *sherpa_shim_create_offline_recognizer(void *cfg);
|
||||
|
||||
// --- Online recognizer config (streaming zipformer transducer) ---
|
||||
void *sherpa_shim_online_recog_config_new(void);
|
||||
void sherpa_shim_online_recog_config_free(void *cfg);
|
||||
void sherpa_shim_online_recog_config_set_transducer_encoder(void *cfg, const char *v);
|
||||
void sherpa_shim_online_recog_config_set_transducer_decoder(void *cfg, const char *v);
|
||||
void sherpa_shim_online_recog_config_set_transducer_joiner(void *cfg, const char *v);
|
||||
void sherpa_shim_online_recog_config_set_tokens(void *cfg, const char *v);
|
||||
void sherpa_shim_online_recog_config_set_num_threads(void *cfg, int32_t v);
|
||||
void sherpa_shim_online_recog_config_set_debug(void *cfg, int32_t v);
|
||||
void sherpa_shim_online_recog_config_set_provider(void *cfg, const char *v);
|
||||
void sherpa_shim_online_recog_config_set_feat_sample_rate(void *cfg, int32_t v);
|
||||
void sherpa_shim_online_recog_config_set_feat_feature_dim(void *cfg, int32_t v);
|
||||
void sherpa_shim_online_recog_config_set_decoding_method(void *cfg, const char *v);
|
||||
void sherpa_shim_online_recog_config_set_enable_endpoint(void *cfg, int32_t v);
|
||||
void sherpa_shim_online_recog_config_set_rule1_min_trailing_silence(void *cfg, float v);
|
||||
void sherpa_shim_online_recog_config_set_rule2_min_trailing_silence(void *cfg, float v);
|
||||
void sherpa_shim_online_recog_config_set_rule3_min_utterance_length(void *cfg, float v);
|
||||
void *sherpa_shim_create_online_recognizer(void *cfg);
|
||||
|
||||
// --- Result accessors (sherpa-allocated; caller destroys via sherpa's own Destroy*) ---
|
||||
int32_t sherpa_shim_wave_sample_rate(const void *wave);
|
||||
int32_t sherpa_shim_wave_num_samples(const void *wave);
|
||||
const float *sherpa_shim_wave_samples(const void *wave);
|
||||
|
||||
const char *sherpa_shim_offline_result_text(const void *result);
|
||||
const char *sherpa_shim_online_result_text(const void *result);
|
||||
|
||||
int32_t sherpa_shim_generated_audio_sample_rate(const void *audio);
|
||||
int32_t sherpa_shim_generated_audio_n(const void *audio);
|
||||
const float *sherpa_shim_generated_audio_samples(const void *audio);
|
||||
|
||||
int32_t sherpa_shim_speech_segment_start(const void *seg);
|
||||
int32_t sherpa_shim_speech_segment_n(const void *seg);
|
||||
|
||||
// --- TTS streaming callback trampoline -----------------------------
|
||||
// Replaces the //export sherpaTtsGoCallback + callbacks.c bridge pattern.
|
||||
// `callback_ptr` is the C-callable function pointer returned by
|
||||
// purego.NewCallback. `user_data` is an integer the Go side uses to
|
||||
// look up its state (sync.Map keyed by uint64).
|
||||
//
|
||||
// Returns the sherpa-allocated SherpaOnnxGeneratedAudio. Destroy with
|
||||
// SherpaOnnxDestroyOfflineTtsGeneratedAudio (callable directly from
|
||||
// Go via purego).
|
||||
void *sherpa_shim_tts_generate_with_callback(
|
||||
void *tts, const char *text, int32_t sid, float speed,
|
||||
uintptr_t callback_ptr, uintptr_t user_data);
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
||||
#endif
|
||||
23
backend/go/sherpa-onnx/main.go
Normal file
23
backend/go/sherpa-onnx/main.go
Normal file
@@ -0,0 +1,23 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"flag"
|
||||
|
||||
grpc "github.com/mudler/LocalAI/pkg/grpc"
|
||||
)
|
||||
|
||||
var (
|
||||
addr = flag.String("addr", "localhost:50051", "the address to connect to")
|
||||
)
|
||||
|
||||
func main() {
|
||||
flag.Parse()
|
||||
|
||||
if err := loadSherpaLibs(); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
if err := grpc.StartServer(*addr, &SherpaBackend{}); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
51
backend/go/sherpa-onnx/package.sh
Executable file
51
backend/go/sherpa-onnx/package.sh
Executable file
@@ -0,0 +1,51 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
CURDIR=$(dirname "$(realpath $0)")
|
||||
REPO_ROOT="${CURDIR}/../../.."
|
||||
|
||||
mkdir -p $CURDIR/package/lib
|
||||
|
||||
cp -avf $CURDIR/sherpa-onnx $CURDIR/package/
|
||||
cp -avf $CURDIR/run.sh $CURDIR/package/
|
||||
cp -rfLv $CURDIR/backend-assets/lib/* $CURDIR/package/lib/
|
||||
|
||||
if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then
|
||||
echo "Detected x86_64 architecture, copying x86_64 libraries..."
|
||||
cp -arfLv /lib64/ld-linux-x86-64.so.2 $CURDIR/package/lib/ld.so
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libc.so.6 $CURDIR/package/lib/libc.so.6
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libgcc_s.so.1 $CURDIR/package/lib/libgcc_s.so.1
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libstdc++.so.6 $CURDIR/package/lib/libstdc++.so.6
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libm.so.6 $CURDIR/package/lib/libm.so.6
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libgomp.so.1 $CURDIR/package/lib/libgomp.so.1
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libdl.so.2 $CURDIR/package/lib/libdl.so.2
|
||||
cp -arfLv /lib/x86_64-linux-gnu/librt.so.1 $CURDIR/package/lib/librt.so.1
|
||||
cp -arfLv /lib/x86_64-linux-gnu/libpthread.so.0 $CURDIR/package/lib/libpthread.so.0
|
||||
elif [ -f "/lib/ld-linux-aarch64.so.1" ]; then
|
||||
echo "Detected ARM64 architecture, copying ARM64 libraries..."
|
||||
cp -arfLv /lib/ld-linux-aarch64.so.1 $CURDIR/package/lib/ld.so
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libc.so.6 $CURDIR/package/lib/libc.so.6
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libgcc_s.so.1 $CURDIR/package/lib/libgcc_s.so.1
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libstdc++.so.6 $CURDIR/package/lib/libstdc++.so.6
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libm.so.6 $CURDIR/package/lib/libm.so.6
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libgomp.so.1 $CURDIR/package/lib/libgomp.so.1
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libdl.so.2 $CURDIR/package/lib/libdl.so.2
|
||||
cp -arfLv /lib/aarch64-linux-gnu/librt.so.1 $CURDIR/package/lib/librt.so.1
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libpthread.so.0 $CURDIR/package/lib/libpthread.so.0
|
||||
elif [ $(uname -s) = "Darwin" ]; then
|
||||
echo "Detected Darwin"
|
||||
else
|
||||
echo "Error: Could not detect architecture"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
GPU_LIB_SCRIPT="${REPO_ROOT}/scripts/build/package-gpu-libs.sh"
|
||||
if [ -f "$GPU_LIB_SCRIPT" ]; then
|
||||
echo "Packaging GPU libraries for BUILD_TYPE=${BUILD_TYPE:-cpu}..."
|
||||
source "$GPU_LIB_SCRIPT" "$CURDIR/package/lib"
|
||||
package_gpu_libs
|
||||
fi
|
||||
|
||||
echo "Packaging completed successfully"
|
||||
ls -liah $CURDIR/package/
|
||||
ls -liah $CURDIR/package/lib/
|
||||
13
backend/go/sherpa-onnx/run.sh
Executable file
13
backend/go/sherpa-onnx/run.sh
Executable file
@@ -0,0 +1,13 @@
|
||||
#!/bin/bash
|
||||
set -ex
|
||||
|
||||
CURDIR=$(dirname "$(realpath $0)")
|
||||
|
||||
export LD_LIBRARY_PATH=$CURDIR/lib:$LD_LIBRARY_PATH
|
||||
|
||||
if [ -f $CURDIR/lib/ld.so ]; then
|
||||
echo "Using lib/ld.so"
|
||||
exec $CURDIR/lib/ld.so $CURDIR/sherpa-onnx "$@"
|
||||
fi
|
||||
|
||||
exec $CURDIR/sherpa-onnx "$@"
|
||||
12
backend/go/sherpa-onnx/test.sh
Executable file
12
backend/go/sherpa-onnx/test.sh
Executable file
@@ -0,0 +1,12 @@
|
||||
#!/bin/bash
|
||||
# Unit tests for the sherpa-onnx backend. Exercises error-path and
|
||||
# dispatch logic via SherpaBackend directly (no gRPC). Integration
|
||||
# coverage (gRPC TTS / streaming ASR / realtime pipeline) lives in
|
||||
# tests/e2e-backends and tests/e2e and runs against the Docker image.
|
||||
set -e
|
||||
|
||||
CURDIR=$(dirname "$(realpath $0)")
|
||||
cd "$CURDIR"
|
||||
|
||||
PACKAGES=$(go list ./... | grep -v /sources/)
|
||||
go test -v -timeout 60s $PACKAGES
|
||||
@@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# stablediffusion.cpp (ggml)
|
||||
STABLEDIFFUSION_GGML_REPO?=https://github.com/leejet/stable-diffusion.cpp
|
||||
STABLEDIFFUSION_GGML_VERSION?=6b675a5ede9b0edf0a0f44191e8b79d7ef27615a
|
||||
STABLEDIFFUSION_GGML_VERSION?=c97702e1057c2fe13a7074cd9069cb9dd6edc1bf
|
||||
|
||||
CMAKE_ARGS+=-DGGML_MAX_NAME=128
|
||||
|
||||
|
||||
@@ -26,6 +26,10 @@
|
||||
#include "stb_image_resize.h"
|
||||
#include <stdlib.h>
|
||||
#include <regex>
|
||||
#include <errno.h>
|
||||
#include <signal.h>
|
||||
#include <unistd.h>
|
||||
#include <sys/wait.h>
|
||||
|
||||
|
||||
|
||||
@@ -980,6 +984,256 @@ int gen_image(sd_img_gen_params_t *p, int steps, char *dst, float cfg_scale, cha
|
||||
return !ret;
|
||||
}
|
||||
|
||||
// ---------------- Video generation ----------------
|
||||
|
||||
sd_vid_gen_params_t* sd_vid_gen_params_new(void) {
|
||||
sd_vid_gen_params_t *params = (sd_vid_gen_params_t *)std::malloc(sizeof(sd_vid_gen_params_t));
|
||||
sd_vid_gen_params_init(params);
|
||||
sd_sample_params_init(¶ms->sample_params);
|
||||
sd_sample_params_init(¶ms->high_noise_sample_params);
|
||||
sd_cache_params_init(¶ms->cache);
|
||||
return params;
|
||||
}
|
||||
|
||||
// Persistent storage for cleaned video prompts (kept alive for the duration of generation)
|
||||
static std::string cleaned_vid_prompt_storage;
|
||||
static std::string cleaned_vid_negative_prompt_storage;
|
||||
|
||||
void sd_vid_gen_params_set_prompts(sd_vid_gen_params_t *params, const char *prompt, const char *negative_prompt) {
|
||||
lora_vec.clear();
|
||||
lora_strings.clear();
|
||||
|
||||
std::string prompt_str = prompt ? prompt : "";
|
||||
std::string negative_prompt_str = negative_prompt ? negative_prompt : "";
|
||||
|
||||
const char* lora_dir_to_use = lora_dir_path.empty() ? nullptr : lora_dir_path.c_str();
|
||||
|
||||
auto [loras, cleaned_prompt] = parse_loras_from_prompt(prompt_str, lora_dir_to_use);
|
||||
lora_vec = loras;
|
||||
cleaned_vid_prompt_storage = cleaned_prompt;
|
||||
|
||||
auto [neg_loras, cleaned_negative] = parse_loras_from_prompt(negative_prompt_str, lora_dir_to_use);
|
||||
cleaned_vid_negative_prompt_storage = cleaned_negative;
|
||||
|
||||
params->prompt = cleaned_vid_prompt_storage.c_str();
|
||||
params->negative_prompt = cleaned_vid_negative_prompt_storage.c_str();
|
||||
params->loras = lora_vec.empty() ? nullptr : lora_vec.data();
|
||||
params->lora_count = static_cast<uint32_t>(lora_vec.size());
|
||||
}
|
||||
|
||||
void sd_vid_gen_params_set_dimensions(sd_vid_gen_params_t *params, int width, int height) {
|
||||
params->width = width;
|
||||
params->height = height;
|
||||
}
|
||||
|
||||
void sd_vid_gen_params_set_seed(sd_vid_gen_params_t *params, int64_t seed) {
|
||||
params->seed = seed;
|
||||
}
|
||||
|
||||
void sd_vid_gen_params_set_video_frames(sd_vid_gen_params_t *params, int n) {
|
||||
params->video_frames = n;
|
||||
}
|
||||
|
||||
// Load an image file into an sd_image_t, resizing to target dims if needed.
|
||||
// Returns a heap-allocated buffer the caller must free (or nullptr on failure).
|
||||
static uint8_t* load_and_resize_image(const char* path, int target_width, int target_height, sd_image_t* out) {
|
||||
if (!path || strlen(path) == 0) {
|
||||
*out = {0, 0, 0, nullptr};
|
||||
return nullptr;
|
||||
}
|
||||
int c = 0, img_w = 0, img_h = 0;
|
||||
uint8_t* buf = stbi_load(path, &img_w, &img_h, &c, 3);
|
||||
if (!buf) {
|
||||
fprintf(stderr, "Failed to load image from '%s'\n", path);
|
||||
*out = {0, 0, 0, nullptr};
|
||||
return nullptr;
|
||||
}
|
||||
if (img_w != target_width || img_h != target_height) {
|
||||
fprintf(stderr, "Resizing image from %dx%d to %dx%d\n", img_w, img_h, target_width, target_height);
|
||||
uint8_t* resized = (uint8_t*)malloc((size_t)target_width * target_height * 3);
|
||||
if (!resized) { free(buf); *out = {0, 0, 0, nullptr}; return nullptr; }
|
||||
stbir_resize(buf, img_w, img_h, 0,
|
||||
resized, target_width, target_height, 0, STBIR_TYPE_UINT8,
|
||||
3, STBIR_ALPHA_CHANNEL_NONE, 0,
|
||||
STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP,
|
||||
STBIR_FILTER_BOX, STBIR_FILTER_BOX,
|
||||
STBIR_COLORSPACE_SRGB, nullptr);
|
||||
free(buf);
|
||||
buf = resized;
|
||||
}
|
||||
*out = {(uint32_t)target_width, (uint32_t)target_height, 3, buf};
|
||||
return buf;
|
||||
}
|
||||
|
||||
// Pipe raw RGB/RGBA frames to ffmpeg stdin and let it produce an MP4 at dst.
|
||||
// Uses fork+execvp to avoid shell interpretation of dst.
|
||||
static int ffmpeg_mux_raw_to_mp4(sd_image_t* frames, int num_frames, int fps, const char* dst) {
|
||||
if (num_frames <= 0 || !frames || !frames[0].data) {
|
||||
fprintf(stderr, "ffmpeg_mux: empty frames\n");
|
||||
return 1;
|
||||
}
|
||||
int width = (int)frames[0].width;
|
||||
int height = (int)frames[0].height;
|
||||
int channels = (int)frames[0].channel;
|
||||
const char* pix_fmt_in = (channels == 4) ? "rgba" : "rgb24";
|
||||
|
||||
char size_str[32];
|
||||
char fps_str[32];
|
||||
snprintf(size_str, sizeof(size_str), "%dx%d", width, height);
|
||||
snprintf(fps_str, sizeof(fps_str), "%d", fps);
|
||||
|
||||
int pipefd[2];
|
||||
if (pipe(pipefd) != 0) { perror("pipe"); return 1; }
|
||||
|
||||
pid_t pid = fork();
|
||||
if (pid < 0) { perror("fork"); close(pipefd[0]); close(pipefd[1]); return 1; }
|
||||
|
||||
if (pid == 0) {
|
||||
// child
|
||||
close(pipefd[1]);
|
||||
if (dup2(pipefd[0], STDIN_FILENO) < 0) { perror("dup2"); _exit(127); }
|
||||
close(pipefd[0]);
|
||||
std::vector<char*> argv = {
|
||||
const_cast<char*>("ffmpeg"),
|
||||
const_cast<char*>("-y"),
|
||||
const_cast<char*>("-hide_banner"),
|
||||
const_cast<char*>("-loglevel"), const_cast<char*>("warning"),
|
||||
const_cast<char*>("-f"), const_cast<char*>("rawvideo"),
|
||||
const_cast<char*>("-pix_fmt"), const_cast<char*>(pix_fmt_in),
|
||||
const_cast<char*>("-s"), size_str,
|
||||
const_cast<char*>("-framerate"), fps_str,
|
||||
const_cast<char*>("-i"), const_cast<char*>("-"),
|
||||
const_cast<char*>("-c:v"), const_cast<char*>("libx264"),
|
||||
const_cast<char*>("-pix_fmt"), const_cast<char*>("yuv420p"),
|
||||
const_cast<char*>("-movflags"), const_cast<char*>("+faststart"),
|
||||
// Force MP4 container. Distributed LocalAI hands us a staging
|
||||
// path (e.g. /staging/localai-output-NNN.tmp) with a non-standard
|
||||
// extension; relying on filename suffix makes ffmpeg bail with
|
||||
// "Unable to choose an output format".
|
||||
const_cast<char*>("-f"), const_cast<char*>("mp4"),
|
||||
const_cast<char*>(dst),
|
||||
nullptr
|
||||
};
|
||||
execvp(argv[0], argv.data());
|
||||
perror("execvp ffmpeg");
|
||||
_exit(127);
|
||||
}
|
||||
|
||||
// parent
|
||||
close(pipefd[0]);
|
||||
|
||||
// Ignore SIGPIPE so a dying ffmpeg surfaces via write() errno instead of killing us.
|
||||
signal(SIGPIPE, SIG_IGN);
|
||||
|
||||
for (int i = 0; i < num_frames; i++) {
|
||||
if (!frames[i].data) continue;
|
||||
size_t frame_bytes = (size_t)frames[i].width * frames[i].height * frames[i].channel;
|
||||
const uint8_t* p = frames[i].data;
|
||||
size_t remaining = frame_bytes;
|
||||
while (remaining > 0) {
|
||||
ssize_t n = write(pipefd[1], p, remaining);
|
||||
if (n < 0) {
|
||||
if (errno == EINTR) continue;
|
||||
perror("write frame to ffmpeg");
|
||||
close(pipefd[1]);
|
||||
int status;
|
||||
waitpid(pid, &status, 0);
|
||||
return 1;
|
||||
}
|
||||
p += n;
|
||||
remaining -= (size_t)n;
|
||||
}
|
||||
}
|
||||
close(pipefd[1]);
|
||||
|
||||
int status = 0;
|
||||
while (waitpid(pid, &status, 0) < 0) {
|
||||
if (errno != EINTR) { perror("waitpid"); return 1; }
|
||||
}
|
||||
if (!WIFEXITED(status) || WEXITSTATUS(status) != 0) {
|
||||
fprintf(stderr, "ffmpeg exited with status %d\n", status);
|
||||
return 1;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
int gen_video(sd_vid_gen_params_t *p, int steps, char *dst, float cfg_scale, int fps, char *init_image, char *end_image) {
|
||||
if (!p) return 1;
|
||||
if (!dst || strlen(dst) == 0) {
|
||||
fprintf(stderr, "gen_video: dst is empty\n");
|
||||
std::free(p);
|
||||
return 1;
|
||||
}
|
||||
|
||||
std::vector<int> skip_layers = {7, 8, 9};
|
||||
|
||||
fprintf(stderr, "Generating video: %dx%d, frames=%d, fps=%d, steps=%d, cfg=%.2f\n",
|
||||
p->width, p->height, p->video_frames, fps, steps, cfg_scale);
|
||||
|
||||
// Sample params (shared by both low and high-noise passes — MoE models use the high-noise
|
||||
// set during the first phase; single-model Wan2.1 ignores it. Same defaults for both is fine.)
|
||||
p->sample_params.guidance.txt_cfg = cfg_scale;
|
||||
p->sample_params.guidance.slg.layers = skip_layers.data();
|
||||
p->sample_params.guidance.slg.layer_count = skip_layers.size();
|
||||
p->sample_params.sample_method = sample_method;
|
||||
p->sample_params.sample_steps = steps;
|
||||
p->sample_params.scheduler = scheduler;
|
||||
p->sample_params.flow_shift = flow_shift;
|
||||
|
||||
p->high_noise_sample_params.guidance.txt_cfg = cfg_scale;
|
||||
p->high_noise_sample_params.guidance.slg.layers = skip_layers.data();
|
||||
p->high_noise_sample_params.guidance.slg.layer_count = skip_layers.size();
|
||||
p->high_noise_sample_params.sample_method = sample_method;
|
||||
p->high_noise_sample_params.sample_steps = steps;
|
||||
p->high_noise_sample_params.scheduler = scheduler;
|
||||
p->high_noise_sample_params.flow_shift = flow_shift;
|
||||
|
||||
// Load init/end reference images if provided (resized to output dims).
|
||||
uint8_t* init_buf = nullptr;
|
||||
uint8_t* end_buf = nullptr;
|
||||
sd_image_t init_img = {0, 0, 0, nullptr};
|
||||
sd_image_t end_img = {0, 0, 0, nullptr};
|
||||
if (init_image && strlen(init_image) > 0) {
|
||||
init_buf = load_and_resize_image(init_image, p->width, p->height, &init_img);
|
||||
if (!init_buf) { std::free(p); return 1; }
|
||||
}
|
||||
if (end_image && strlen(end_image) > 0) {
|
||||
end_buf = load_and_resize_image(end_image, p->width, p->height, &end_img);
|
||||
if (!end_buf) { if (init_buf) free(init_buf); std::free(p); return 1; }
|
||||
}
|
||||
p->init_image = init_img;
|
||||
p->end_image = end_img;
|
||||
|
||||
// Generate
|
||||
int num_frames_out = 0;
|
||||
sd_image_t* frames = generate_video(sd_c, p, &num_frames_out);
|
||||
std::free(p);
|
||||
|
||||
if (!frames || num_frames_out == 0) {
|
||||
fprintf(stderr, "generate_video produced no frames\n");
|
||||
if (init_buf) free(init_buf);
|
||||
if (end_buf) free(end_buf);
|
||||
return 1;
|
||||
}
|
||||
|
||||
fprintf(stderr, "Generated %d frames, muxing to %s via ffmpeg\n", num_frames_out, dst);
|
||||
|
||||
int rc = ffmpeg_mux_raw_to_mp4(frames, num_frames_out, fps, dst);
|
||||
|
||||
for (int i = 0; i < num_frames_out; i++) {
|
||||
if (frames[i].data) free(frames[i].data);
|
||||
}
|
||||
free(frames);
|
||||
if (init_buf) free(init_buf);
|
||||
if (end_buf) free(end_buf);
|
||||
|
||||
if (rc == 0) {
|
||||
fprintf(stderr, "gen_video done: %s\n", dst);
|
||||
}
|
||||
fflush(stderr);
|
||||
return rc;
|
||||
}
|
||||
|
||||
int unload() {
|
||||
free_sd_ctx(sd_c);
|
||||
return 0;
|
||||
|
||||
@@ -23,6 +23,7 @@ type SDGGML struct {
|
||||
var (
|
||||
LoadModel func(model, model_apth string, options []uintptr, threads int32, diff int) int
|
||||
GenImage func(params uintptr, steps int, dst string, cfgScale float32, srcImage string, strength float32, maskImage string, refImages []uintptr, refImagesCount int) int
|
||||
GenVideo func(params uintptr, steps int, dst string, cfgScale float32, fps int, initImage string, endImage string) int
|
||||
|
||||
TilingParamsSetEnabled func(params uintptr, enabled bool)
|
||||
TilingParamsSetTileSizes func(params uintptr, tileSizeX int, tileSizeY int)
|
||||
@@ -34,6 +35,12 @@ var (
|
||||
ImgGenParamsSetDimensions func(params uintptr, width int, height int)
|
||||
ImgGenParamsSetSeed func(params uintptr, seed int64)
|
||||
ImgGenParamsGetVaeTilingParams func(params uintptr) uintptr
|
||||
|
||||
VidGenParamsNew func() uintptr
|
||||
VidGenParamsSetPrompts func(params uintptr, prompt string, negativePrompt string)
|
||||
VidGenParamsSetDimensions func(params uintptr, width int, height int)
|
||||
VidGenParamsSetSeed func(params uintptr, seed int64)
|
||||
VidGenParamsSetVideoFrames func(params uintptr, n int)
|
||||
)
|
||||
|
||||
// Copied from Purego internal/strings
|
||||
@@ -153,3 +160,58 @@ func (sd *SDGGML) GenerateImage(opts *pb.GenerateImageRequest) error {
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (sd *SDGGML) GenerateVideo(opts *pb.GenerateVideoRequest) error {
|
||||
dst := opts.Dst
|
||||
if dst == "" {
|
||||
return fmt.Errorf("dst is empty")
|
||||
}
|
||||
|
||||
width := int(opts.Width)
|
||||
height := int(opts.Height)
|
||||
if width == 0 {
|
||||
width = 512
|
||||
}
|
||||
if height == 0 {
|
||||
height = 512
|
||||
}
|
||||
|
||||
numFrames := int(opts.NumFrames)
|
||||
if numFrames <= 0 {
|
||||
numFrames = 16
|
||||
}
|
||||
|
||||
fps := int(opts.Fps)
|
||||
if fps <= 0 {
|
||||
fps = 16
|
||||
}
|
||||
|
||||
steps := int(opts.Step)
|
||||
if steps <= 0 {
|
||||
steps = 20
|
||||
}
|
||||
|
||||
cfg := opts.CfgScale
|
||||
if cfg == 0 {
|
||||
cfg = sd.cfgScale
|
||||
}
|
||||
if cfg == 0 {
|
||||
cfg = 5.0
|
||||
}
|
||||
|
||||
// sd_vid_gen_params_new allocates; gen_video frees it after the generation call.
|
||||
p := VidGenParamsNew()
|
||||
VidGenParamsSetPrompts(p, opts.Prompt, opts.NegativePrompt)
|
||||
VidGenParamsSetDimensions(p, width, height)
|
||||
VidGenParamsSetSeed(p, int64(opts.Seed))
|
||||
VidGenParamsSetVideoFrames(p, numFrames)
|
||||
|
||||
fmt.Fprintf(os.Stderr, "GenerateVideo: dst=%s size=%dx%d frames=%d fps=%d steps=%d cfg=%.2f\n",
|
||||
dst, width, height, numFrames, fps, steps, cfg)
|
||||
|
||||
ret := GenVideo(p, steps, dst, cfg, fps, opts.StartImage, opts.EndImage)
|
||||
if ret != 0 {
|
||||
return fmt.Errorf("video inference failed (code %d)", ret)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -18,6 +18,13 @@ void sd_img_gen_params_set_seed(sd_img_gen_params_t *params, int64_t seed);
|
||||
|
||||
int load_model(const char *model, char *model_path, char* options[], int threads, int diffusionModel);
|
||||
int gen_image(sd_img_gen_params_t *p, int steps, char *dst, float cfg_scale, char *src_image, float strength, char *mask_image, char* ref_images[], int ref_images_count);
|
||||
|
||||
sd_vid_gen_params_t* sd_vid_gen_params_new(void);
|
||||
void sd_vid_gen_params_set_prompts(sd_vid_gen_params_t *params, const char *prompt, const char *negative_prompt);
|
||||
void sd_vid_gen_params_set_dimensions(sd_vid_gen_params_t *params, int width, int height);
|
||||
void sd_vid_gen_params_set_seed(sd_vid_gen_params_t *params, int64_t seed);
|
||||
void sd_vid_gen_params_set_video_frames(sd_vid_gen_params_t *params, int n);
|
||||
int gen_video(sd_vid_gen_params_t *p, int steps, char *dst, float cfg_scale, int fps, char *init_image, char *end_image);
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
||||
@@ -32,6 +32,7 @@ func main() {
|
||||
libFuncs := []LibFuncs{
|
||||
{&LoadModel, "load_model"},
|
||||
{&GenImage, "gen_image"},
|
||||
{&GenVideo, "gen_video"},
|
||||
{&TilingParamsSetEnabled, "sd_tiling_params_set_enabled"},
|
||||
{&TilingParamsSetTileSizes, "sd_tiling_params_set_tile_sizes"},
|
||||
{&TilingParamsSetRelSizes, "sd_tiling_params_set_rel_sizes"},
|
||||
@@ -42,6 +43,12 @@ func main() {
|
||||
{&ImgGenParamsSetDimensions, "sd_img_gen_params_set_dimensions"},
|
||||
{&ImgGenParamsSetSeed, "sd_img_gen_params_set_seed"},
|
||||
{&ImgGenParamsGetVaeTilingParams, "sd_img_gen_params_get_vae_tiling_params"},
|
||||
|
||||
{&VidGenParamsNew, "sd_vid_gen_params_new"},
|
||||
{&VidGenParamsSetPrompts, "sd_vid_gen_params_set_prompts"},
|
||||
{&VidGenParamsSetDimensions, "sd_vid_gen_params_set_dimensions"},
|
||||
{&VidGenParamsSetSeed, "sd_vid_gen_params_set_seed"},
|
||||
{&VidGenParamsSetVideoFrames, "sd_vid_gen_params_set_video_frames"},
|
||||
}
|
||||
|
||||
for _, lf := range libFuncs {
|
||||
|
||||
@@ -56,5 +56,6 @@ func (v *Voxtral) AudioTranscription(opts *pb.TranscriptRequest) (pb.TranscriptR
|
||||
return pb.TranscriptResult{
|
||||
Segments: segments,
|
||||
Text: text,
|
||||
Language: opts.Language,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# whisper.cpp version
|
||||
WHISPER_REPO?=https://github.com/ggml-org/whisper.cpp
|
||||
WHISPER_CPP_VERSION?=95ea8f9bfb03a15db08a8989966fd1ae3361e20d
|
||||
WHISPER_CPP_VERSION?=fc674574ca27cac59a15e5b22a09b9d9ad62aafe
|
||||
SO_TARGET?=libgowhisper.so
|
||||
|
||||
CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF
|
||||
|
||||
@@ -120,6 +120,12 @@ func (w *Whisper) AudioTranscription(opts *pb.TranscriptRequest) (pb.TranscriptR
|
||||
}
|
||||
|
||||
data := buf.AsFloat32Buffer().Data
|
||||
// whisper.cpp resamples to 16 kHz internally; this matches buf.Format.SampleRate
|
||||
// for the converted file produced by AudioToWav above.
|
||||
var duration float32
|
||||
if buf.Format != nil && buf.Format.SampleRate > 0 {
|
||||
duration = float32(len(data)) / float32(buf.Format.SampleRate)
|
||||
}
|
||||
segsLen := uintptr(0xdeadbeef)
|
||||
segsLenPtr := unsafe.Pointer(&segsLen)
|
||||
|
||||
@@ -158,5 +164,7 @@ func (w *Whisper) AudioTranscription(opts *pb.TranscriptRequest) (pb.TranscriptR
|
||||
return pb.TranscriptResult{
|
||||
Segments: segments,
|
||||
Text: strings.TrimSpace(text),
|
||||
Language: opts.Language,
|
||||
Duration: duration,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -43,6 +43,35 @@
|
||||
- CPU
|
||||
capabilities:
|
||||
default: "cpu-ik-llama-cpp"
|
||||
- &turboquant
|
||||
name: "turboquant"
|
||||
alias: "turboquant"
|
||||
license: mit
|
||||
description: |
|
||||
Fork of llama.cpp adding the TurboQuant KV-cache quantization scheme.
|
||||
Reuses the LocalAI llama.cpp gRPC server sources against the fork's libllama.
|
||||
urls:
|
||||
- https://github.com/TheTom/llama-cpp-turboquant
|
||||
tags:
|
||||
- text-to-text
|
||||
- LLM
|
||||
- CPU
|
||||
- GPU
|
||||
- CUDA
|
||||
- HIP
|
||||
- turboquant
|
||||
- kv-cache
|
||||
capabilities:
|
||||
default: "cpu-turboquant"
|
||||
nvidia: "cuda12-turboquant"
|
||||
intel: "intel-sycl-f16-turboquant"
|
||||
amd: "rocm-turboquant"
|
||||
vulkan: "vulkan-turboquant"
|
||||
nvidia-l4t: "nvidia-l4t-arm64-turboquant"
|
||||
nvidia-cuda-13: "cuda13-turboquant"
|
||||
nvidia-cuda-12: "cuda12-turboquant"
|
||||
nvidia-l4t-cuda-12: "nvidia-l4t-arm64-turboquant"
|
||||
nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-turboquant"
|
||||
- &whispercpp
|
||||
name: "whisper"
|
||||
alias: "whisper"
|
||||
@@ -139,6 +168,43 @@
|
||||
nvidia-cuda-13: "cuda13-rfdetr"
|
||||
nvidia-cuda-12: "cuda12-rfdetr"
|
||||
nvidia-l4t-cuda-12: "nvidia-l4t-arm64-rfdetr"
|
||||
- &insightface
|
||||
name: "insightface"
|
||||
alias: "insightface"
|
||||
# Upstream insightface library is MIT. The pretrained model packs
|
||||
# (buffalo_l, buffalo_s, antelopev2) are released for NON-COMMERCIAL
|
||||
# research use only. The backend image also pre-bakes OpenCV Zoo
|
||||
# YuNet + SFace (Apache 2.0) for commercial use. Pick the engine
|
||||
# via model-gallery entries (insightface-buffalo-l / insightface-opencv
|
||||
# / insightface-buffalo-s) or set `options` in your model YAML.
|
||||
license: "mixed"
|
||||
description: |
|
||||
Face recognition backend powered by `insightface` (ONNX Runtime).
|
||||
Provides face verification (/v1/face/verify), face analysis
|
||||
(/v1/face/analyze), face embedding (/v1/embeddings), face
|
||||
detection (/v1/detection), and 1:N identification
|
||||
(/v1/face/{register,identify,forget}).
|
||||
Ships two engines in a single image: one that drives the insightface
|
||||
model packs (buffalo_l/s/m/sc, antelopev2 — non-commercial research
|
||||
use only) and one that drives OpenCV Zoo's YuNet + SFace pair
|
||||
(Apache 2.0 — commercial-safe). Select via `options: ["engine:..."]`
|
||||
in your model YAML, or install one of the ready-made model-gallery
|
||||
entries under the `insightface-*` prefix.
|
||||
The backend image contains only code and Python deps; all model
|
||||
weights are managed by LocalAI's gallery download mechanism.
|
||||
urls:
|
||||
- https://github.com/deepinsight/insightface
|
||||
- https://github.com/opencv/opencv_zoo
|
||||
tags:
|
||||
- face-recognition
|
||||
- face-verification
|
||||
- face-embedding
|
||||
- gpu
|
||||
- cpu
|
||||
capabilities:
|
||||
default: "cpu-insightface"
|
||||
nvidia: "cuda12-insightface"
|
||||
nvidia-cuda-12: "cuda12-insightface"
|
||||
- &sam3cpp
|
||||
name: "sam3-cpp"
|
||||
alias: "sam3-cpp"
|
||||
@@ -198,6 +264,28 @@
|
||||
intel: "intel-vllm"
|
||||
nvidia-cuda-12: "cuda12-vllm"
|
||||
cpu: "cpu-vllm"
|
||||
- &sglang
|
||||
name: "sglang"
|
||||
license: apache-2.0
|
||||
urls:
|
||||
- https://github.com/sgl-project/sglang
|
||||
tags:
|
||||
- text-to-text
|
||||
- multimodal
|
||||
icon: https://raw.githubusercontent.com/sgl-project/sglang/main/assets/logo.png
|
||||
description: |
|
||||
SGLang is a fast serving framework for large language models and vision language models.
|
||||
It co-designs the backend runtime (RadixAttention, continuous batching, structured
|
||||
decoding) and the frontend language to make interaction with models faster and more
|
||||
controllable. Features include fast backend runtime, flexible frontend language,
|
||||
extensive model support, and an active community.
|
||||
alias: "sglang"
|
||||
capabilities:
|
||||
nvidia: "cuda12-sglang"
|
||||
amd: "rocm-sglang"
|
||||
intel: "intel-sglang"
|
||||
nvidia-cuda-12: "cuda12-sglang"
|
||||
cpu: "cpu-sglang"
|
||||
- &vllm-omni
|
||||
name: "vllm-omni"
|
||||
license: apache-2.0
|
||||
@@ -332,6 +420,34 @@
|
||||
intel: "intel-rerankers"
|
||||
amd: "rocm-rerankers"
|
||||
metal: "metal-rerankers"
|
||||
- &tinygrad
|
||||
name: "tinygrad"
|
||||
alias: "tinygrad"
|
||||
license: MIT
|
||||
description: |
|
||||
tinygrad is a minimalist deep-learning framework with zero runtime
|
||||
dependencies that targets CUDA, ROCm, Metal, WebGPU and CPU (CLANG).
|
||||
The LocalAI tinygrad backend exposes a single multimodal runtime that
|
||||
covers LLM text generation (Llama / Qwen / Mistral via safetensors or
|
||||
GGUF) with native tool-call extraction, BERT-family embeddings,
|
||||
Stable Diffusion 1.x / 2 / XL image generation, and Whisper speech-to-text.
|
||||
|
||||
Single image: tinygrad generates its own GPU kernels and dlopens the
|
||||
host driver libraries at runtime, so there is no per-toolkit build
|
||||
split. The same image runs CPU-only or accelerates against
|
||||
CUDA / ROCm / Metal when the host driver is visible.
|
||||
urls:
|
||||
- https://github.com/tinygrad/tinygrad
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-tinygrad"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-tinygrad
|
||||
tags:
|
||||
- text-to-text
|
||||
- LLM
|
||||
- embeddings
|
||||
- image-generation
|
||||
- transcription
|
||||
- multimodal
|
||||
- &transformers
|
||||
name: "transformers"
|
||||
icon: https://avatars.githubusercontent.com/u/25720743?s=200&v=4
|
||||
@@ -508,7 +624,6 @@
|
||||
alias: "whisperx"
|
||||
capabilities:
|
||||
nvidia: "cuda12-whisperx"
|
||||
amd: "rocm-whisperx"
|
||||
metal: "metal-whisperx"
|
||||
default: "cpu-whisperx"
|
||||
nvidia-cuda-13: "cuda13-whisperx"
|
||||
@@ -891,6 +1006,23 @@
|
||||
nvidia: "cuda12-neutts"
|
||||
amd: "rocm-neutts"
|
||||
nvidia-cuda-12: "cuda12-neutts"
|
||||
- &sherpa-onnx
|
||||
name: "sherpa-onnx"
|
||||
alias: "sherpa-onnx"
|
||||
urls:
|
||||
- https://k2-fsa.github.io/sherpa/onnx/
|
||||
description: |
|
||||
Sherpa-ONNX backend for text-to-speech (VITS, Matcha, Kokoro), speech-to-text (Whisper, Paraformer, SenseVoice, Omnilingual ASR CTC), and voice activity detection via ONNX Runtime.
|
||||
Supports multi-speaker voices, 1600+ language ASR, and GPU acceleration.
|
||||
tags:
|
||||
- text-to-speech
|
||||
- TTS
|
||||
- speech-to-text
|
||||
- ASR
|
||||
capabilities:
|
||||
default: "cpu-sherpa-onnx"
|
||||
nvidia: "cuda12-sherpa-onnx"
|
||||
nvidia-cuda-12: "cuda12-sherpa-onnx"
|
||||
- !!merge <<: *neutts
|
||||
name: "neutts-development"
|
||||
capabilities:
|
||||
@@ -916,6 +1048,33 @@
|
||||
name: "ik-llama-cpp-development"
|
||||
capabilities:
|
||||
default: "cpu-ik-llama-cpp-development"
|
||||
- !!merge <<: *turboquant
|
||||
name: "turboquant-development"
|
||||
capabilities:
|
||||
default: "cpu-turboquant-development"
|
||||
nvidia: "cuda12-turboquant-development"
|
||||
intel: "intel-sycl-f16-turboquant-development"
|
||||
amd: "rocm-turboquant-development"
|
||||
vulkan: "vulkan-turboquant-development"
|
||||
nvidia-l4t: "nvidia-l4t-arm64-turboquant-development"
|
||||
nvidia-cuda-13: "cuda13-turboquant-development"
|
||||
nvidia-cuda-12: "cuda12-turboquant-development"
|
||||
nvidia-l4t-cuda-12: "nvidia-l4t-arm64-turboquant-development"
|
||||
nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-turboquant-development"
|
||||
- !!merge <<: *stablediffusionggml
|
||||
name: "stablediffusion-ggml-development"
|
||||
capabilities:
|
||||
default: "cpu-stablediffusion-ggml-development"
|
||||
nvidia: "cuda12-stablediffusion-ggml-development"
|
||||
intel: "intel-sycl-f16-stablediffusion-ggml-development"
|
||||
# amd: "rocm-stablediffusion-ggml-development"
|
||||
vulkan: "vulkan-stablediffusion-ggml-development"
|
||||
nvidia-l4t: "nvidia-l4t-arm64-stablediffusion-ggml-development"
|
||||
metal: "metal-stablediffusion-ggml-development"
|
||||
nvidia-cuda-13: "cuda13-stablediffusion-ggml-development"
|
||||
nvidia-cuda-12: "cuda12-stablediffusion-ggml-development"
|
||||
nvidia-l4t-cuda-12: "nvidia-l4t-arm64-stablediffusion-ggml-development"
|
||||
nvidia-l4t-cuda-13: "cuda13-nvidia-l4t-arm64-stablediffusion-ggml-development"
|
||||
- !!merge <<: *neutts
|
||||
name: "cpu-neutts"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-neutts"
|
||||
@@ -1357,6 +1516,97 @@
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-cpu-ik-llama-cpp"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-cpu-ik-llama-cpp
|
||||
## turboquant
|
||||
- !!merge <<: *turboquant
|
||||
name: "cpu-turboquant"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-turboquant"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-cpu-turboquant
|
||||
- !!merge <<: *turboquant
|
||||
name: "cpu-turboquant-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-cpu-turboquant"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-cpu-turboquant
|
||||
- !!merge <<: *turboquant
|
||||
name: "cuda12-turboquant"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-12-turboquant"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-nvidia-cuda-12-turboquant
|
||||
- !!merge <<: *turboquant
|
||||
name: "cuda12-turboquant-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-12-turboquant"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-nvidia-cuda-12-turboquant
|
||||
- !!merge <<: *turboquant
|
||||
name: "cuda13-turboquant"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-13-turboquant"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-nvidia-cuda-13-turboquant
|
||||
- !!merge <<: *turboquant
|
||||
name: "cuda13-turboquant-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-13-turboquant"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-nvidia-cuda-13-turboquant
|
||||
- !!merge <<: *turboquant
|
||||
name: "rocm-turboquant"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-rocm-hipblas-turboquant"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-rocm-hipblas-turboquant
|
||||
- !!merge <<: *turboquant
|
||||
name: "rocm-turboquant-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-rocm-hipblas-turboquant"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-rocm-hipblas-turboquant
|
||||
- !!merge <<: *turboquant
|
||||
name: "intel-sycl-f32-turboquant"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-intel-sycl-f32-turboquant"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-intel-sycl-f32-turboquant
|
||||
- !!merge <<: *turboquant
|
||||
name: "intel-sycl-f32-turboquant-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-intel-sycl-f32-turboquant"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-intel-sycl-f32-turboquant
|
||||
- !!merge <<: *turboquant
|
||||
name: "intel-sycl-f16-turboquant"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-intel-sycl-f16-turboquant"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-intel-sycl-f16-turboquant
|
||||
- !!merge <<: *turboquant
|
||||
name: "intel-sycl-f16-turboquant-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-intel-sycl-f16-turboquant"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-intel-sycl-f16-turboquant
|
||||
- !!merge <<: *turboquant
|
||||
name: "vulkan-turboquant"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-vulkan-turboquant"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-vulkan-turboquant
|
||||
- !!merge <<: *turboquant
|
||||
name: "vulkan-turboquant-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-vulkan-turboquant"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-vulkan-turboquant
|
||||
- !!merge <<: *turboquant
|
||||
name: "nvidia-l4t-arm64-turboquant"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-nvidia-l4t-arm64-turboquant"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-nvidia-l4t-arm64-turboquant
|
||||
- !!merge <<: *turboquant
|
||||
name: "nvidia-l4t-arm64-turboquant-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-nvidia-l4t-arm64-turboquant"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-nvidia-l4t-arm64-turboquant
|
||||
- !!merge <<: *turboquant
|
||||
name: "cuda13-nvidia-l4t-arm64-turboquant"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-nvidia-l4t-cuda-13-arm64-turboquant"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-nvidia-l4t-cuda-13-arm64-turboquant
|
||||
- !!merge <<: *turboquant
|
||||
name: "cuda13-nvidia-l4t-arm64-turboquant-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-nvidia-l4t-cuda-13-arm64-turboquant"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-nvidia-l4t-cuda-13-arm64-turboquant
|
||||
## whisper
|
||||
- !!merge <<: *whispercpp
|
||||
name: "nvidia-l4t-arm64-whisper"
|
||||
@@ -1605,6 +1855,54 @@
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-cpu-vllm"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-cpu-vllm
|
||||
# sglang
|
||||
- !!merge <<: *sglang
|
||||
name: "sglang-development"
|
||||
capabilities:
|
||||
nvidia: "cuda12-sglang-development"
|
||||
amd: "rocm-sglang-development"
|
||||
intel: "intel-sglang-development"
|
||||
cpu: "cpu-sglang-development"
|
||||
- !!merge <<: *sglang
|
||||
name: "cuda12-sglang"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-12-sglang"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-nvidia-cuda-12-sglang
|
||||
- !!merge <<: *sglang
|
||||
name: "rocm-sglang"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-rocm-hipblas-sglang"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-rocm-hipblas-sglang
|
||||
- !!merge <<: *sglang
|
||||
name: "intel-sglang"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-intel-sglang"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-intel-sglang
|
||||
- !!merge <<: *sglang
|
||||
name: "cpu-sglang"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-sglang"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-cpu-sglang
|
||||
- !!merge <<: *sglang
|
||||
name: "cuda12-sglang-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-12-sglang"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-nvidia-cuda-12-sglang
|
||||
- !!merge <<: *sglang
|
||||
name: "rocm-sglang-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-rocm-hipblas-sglang"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-rocm-hipblas-sglang
|
||||
- !!merge <<: *sglang
|
||||
name: "intel-sglang-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-intel-sglang"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-intel-sglang
|
||||
- !!merge <<: *sglang
|
||||
name: "cpu-sglang-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-cpu-sglang"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-cpu-sglang
|
||||
# vllm-omni
|
||||
- !!merge <<: *vllm-omni
|
||||
name: "vllm-omni-development"
|
||||
@@ -1860,6 +2158,15 @@
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-rerankers"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-metal-darwin-arm64-rerankers
|
||||
## tinygrad
|
||||
## Single image — the meta anchor above carries the latest uri directly
|
||||
## since there is only one variant. The development entry below points at
|
||||
## the master tag.
|
||||
- !!merge <<: *tinygrad
|
||||
name: "tinygrad-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-tinygrad"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-tinygrad
|
||||
## Transformers
|
||||
- !!merge <<: *transformers
|
||||
name: "transformers-development"
|
||||
@@ -2491,7 +2798,6 @@
|
||||
name: "whisperx-development"
|
||||
capabilities:
|
||||
nvidia: "cuda12-whisperx-development"
|
||||
amd: "rocm-whisperx-development"
|
||||
metal: "metal-whisperx-development"
|
||||
default: "cpu-whisperx-development"
|
||||
nvidia-cuda-13: "cuda13-whisperx-development"
|
||||
@@ -2517,16 +2823,6 @@
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-12-whisperx"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-nvidia-cuda-12-whisperx
|
||||
- !!merge <<: *whisperx
|
||||
name: "rocm-whisperx"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-rocm-hipblas-whisperx"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-rocm-hipblas-whisperx
|
||||
- !!merge <<: *whisperx
|
||||
name: "rocm-whisperx-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-rocm-hipblas-whisperx"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-rocm-hipblas-whisperx
|
||||
- !!merge <<: *whisperx
|
||||
name: "cuda13-whisperx"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-13-whisperx"
|
||||
@@ -3467,3 +3763,118 @@
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-metal-darwin-arm64-llama-cpp-quantization"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-metal-darwin-arm64-llama-cpp-quantization
|
||||
# insightface (face recognition) — development and concrete image entries
|
||||
- !!merge <<: *insightface
|
||||
name: "insightface-development"
|
||||
capabilities:
|
||||
default: "cpu-insightface-development"
|
||||
nvidia: "cuda12-insightface-development"
|
||||
nvidia-cuda-12: "cuda12-insightface-development"
|
||||
- !!merge <<: *insightface
|
||||
name: "cpu-insightface"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-insightface"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-cpu-insightface
|
||||
- !!merge <<: *insightface
|
||||
name: "cuda12-insightface"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-12-insightface"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-nvidia-cuda-12-insightface
|
||||
- !!merge <<: *insightface
|
||||
name: "cpu-insightface-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-cpu-insightface"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-cpu-insightface
|
||||
- !!merge <<: *insightface
|
||||
name: "cuda12-insightface-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-12-insightface"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-nvidia-cuda-12-insightface
|
||||
|
||||
# speaker-recognition (voice/speaker biometrics) — Apache-2.0 stack
|
||||
- &speakerrecognition
|
||||
name: "speaker-recognition"
|
||||
alias: "speaker-recognition"
|
||||
# SpeechBrain is Apache-2.0. WeSpeaker / 3D-Speaker ONNX exports are
|
||||
# Apache-2.0. The backend itself ships only Python deps — all model
|
||||
# weights flow through LocalAI's gallery download mechanism (or
|
||||
# SpeechBrain's built-in HF auto-download at first LoadModel).
|
||||
license: apache-2.0
|
||||
description: |
|
||||
Speaker (voice) recognition backend — the audio analog to
|
||||
insightface. Wraps SpeechBrain ECAPA-TDNN (default engine, 192-d
|
||||
embeddings, ~1.9% EER on VoxCeleb) plus an OnnxDirectEngine for
|
||||
pre-exported WeSpeaker / 3D-Speaker ONNX models.
|
||||
|
||||
Exposes speaker verification (/v1/voice/verify), speaker embedding
|
||||
(/v1/voice/embed), speaker analysis (/v1/voice/analyze), and 1:N
|
||||
speaker identification (/v1/voice/{register,identify,forget}).
|
||||
Registrations use LocalAI's built-in vector store — same in-memory
|
||||
backing the face-recognition registry uses, separate instance.
|
||||
urls:
|
||||
- https://speechbrain.github.io/
|
||||
- https://github.com/wenet-e2e/wespeaker
|
||||
- https://github.com/modelscope/3D-Speaker
|
||||
tags:
|
||||
- voice-recognition
|
||||
- speaker-verification
|
||||
- speaker-embedding
|
||||
- gpu
|
||||
- cpu
|
||||
capabilities:
|
||||
default: "cpu-speaker-recognition"
|
||||
nvidia: "cuda12-speaker-recognition"
|
||||
nvidia-cuda-12: "cuda12-speaker-recognition"
|
||||
- !!merge <<: *speakerrecognition
|
||||
name: "speaker-recognition-development"
|
||||
capabilities:
|
||||
default: "cpu-speaker-recognition-development"
|
||||
nvidia: "cuda12-speaker-recognition-development"
|
||||
nvidia-cuda-12: "cuda12-speaker-recognition-development"
|
||||
- !!merge <<: *speakerrecognition
|
||||
name: "cpu-speaker-recognition"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-speaker-recognition"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-cpu-speaker-recognition
|
||||
- !!merge <<: *speakerrecognition
|
||||
name: "cuda12-speaker-recognition"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-12-speaker-recognition"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-nvidia-cuda-12-speaker-recognition
|
||||
- !!merge <<: *speakerrecognition
|
||||
name: "cpu-speaker-recognition-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-cpu-speaker-recognition"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-cpu-speaker-recognition
|
||||
- !!merge <<: *speakerrecognition
|
||||
name: "cuda12-speaker-recognition-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-12-speaker-recognition"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-nvidia-cuda-12-speaker-recognition
|
||||
## sherpa-onnx
|
||||
- !!merge <<: *sherpa-onnx
|
||||
name: "sherpa-onnx-development"
|
||||
capabilities:
|
||||
default: "cpu-sherpa-onnx-development"
|
||||
nvidia: "cuda12-sherpa-onnx-development"
|
||||
nvidia-cuda-12: "cuda12-sherpa-onnx-development"
|
||||
- !!merge <<: *sherpa-onnx
|
||||
name: "cpu-sherpa-onnx"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-sherpa-onnx"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-cpu-sherpa-onnx
|
||||
- !!merge <<: *sherpa-onnx
|
||||
name: "cpu-sherpa-onnx-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-cpu-sherpa-onnx"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-cpu-sherpa-onnx
|
||||
- !!merge <<: *sherpa-onnx
|
||||
name: "cuda12-sherpa-onnx"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-12-sherpa-onnx"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-nvidia-cuda-12-sherpa-onnx
|
||||
- !!merge <<: *sherpa-onnx
|
||||
name: "cuda12-sherpa-onnx-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-12-sherpa-onnx"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-nvidia-cuda-12-sherpa-onnx
|
||||
|
||||
@@ -344,7 +344,16 @@ function ensureVenv() {
|
||||
|
||||
if [ ! -d "${EDIR}/venv" ]; then
|
||||
if [ "x${USE_PIP}" == "xtrue" ]; then
|
||||
"${interpreter}" -m venv --copies "${EDIR}/venv"
|
||||
# --copies is only needed when we will later relocate the venv via
|
||||
# _makeVenvPortable (PORTABLE_PYTHON=true). Some Python builds —
|
||||
# notably macOS system Python — refuse to create a venv with
|
||||
# --copies because the build doesn't support it. Fall back to
|
||||
# symlinks in that case.
|
||||
local venv_args=""
|
||||
if [ "x${PORTABLE_PYTHON}" == "xtrue" ]; then
|
||||
venv_args="--copies"
|
||||
fi
|
||||
"${interpreter}" -m venv ${venv_args} "${EDIR}/venv"
|
||||
source "${EDIR}/venv/bin/activate"
|
||||
"${interpreter}" -m pip install --upgrade pip
|
||||
else
|
||||
|
||||
100
backend/python/common/mlx_utils.py
Normal file
100
backend/python/common/mlx_utils.py
Normal file
@@ -0,0 +1,100 @@
|
||||
"""Shared utilities for the mlx and mlx-vlm gRPC backends.
|
||||
|
||||
These helpers wrap mlx-lm's and mlx-vlm's native tool-parser modules, which
|
||||
auto-detect the right parser from the model's chat template. Each tool
|
||||
module exposes ``tool_call_start``, ``tool_call_end`` and
|
||||
``parse_tool_call(text, tools) -> dict | list[dict]``.
|
||||
|
||||
The split-reasoning helper is generic enough to work with any think-start /
|
||||
think-end delimiter pair.
|
||||
"""
|
||||
import json
|
||||
import re
|
||||
import sys
|
||||
import uuid
|
||||
|
||||
|
||||
def split_reasoning(text, think_start, think_end):
|
||||
"""Split ``<think>...</think>`` blocks out of ``text``.
|
||||
|
||||
Returns ``(reasoning_content, remaining_text)``. When ``think_start`` is
|
||||
empty or not found, returns ``("", text)`` unchanged.
|
||||
"""
|
||||
if not think_start or not text or think_start not in text:
|
||||
return "", text
|
||||
pattern = re.compile(
|
||||
re.escape(think_start) + r"(.*?)" + re.escape(think_end or ""),
|
||||
re.DOTALL,
|
||||
)
|
||||
reasoning_parts = pattern.findall(text)
|
||||
if not reasoning_parts:
|
||||
return "", text
|
||||
remaining = pattern.sub("", text).strip()
|
||||
return "\n".join(p.strip() for p in reasoning_parts), remaining
|
||||
|
||||
|
||||
def parse_tool_calls(text, tool_module, tools):
|
||||
"""Extract tool calls from ``text`` using a mlx-lm tool module.
|
||||
|
||||
Ports the ``process_tool_calls`` logic from
|
||||
``mlx_vlm/server.py`` (v0.10 onwards). ``tool_module`` must expose
|
||||
``tool_call_start``, ``tool_call_end`` and ``parse_tool_call``.
|
||||
|
||||
Returns ``(calls, remaining_text)`` where ``calls`` is a list of dicts:
|
||||
|
||||
[{"index": int, "id": str, "name": str, "arguments": str (JSON)}]
|
||||
|
||||
and ``remaining_text`` is the free-form text with the tool call blocks
|
||||
removed. ``(calls, text)`` is returned unchanged if ``tool_module`` is
|
||||
``None`` or the start delimiter isn't present.
|
||||
"""
|
||||
if tool_module is None or not text:
|
||||
return [], text
|
||||
start = getattr(tool_module, "tool_call_start", None)
|
||||
end = getattr(tool_module, "tool_call_end", None)
|
||||
parse_fn = getattr(tool_module, "parse_tool_call", None)
|
||||
if not start or parse_fn is None or start not in text:
|
||||
return [], text
|
||||
|
||||
if end == "" or end is None:
|
||||
pattern = re.compile(
|
||||
re.escape(start) + r".*?(?:\n|$)",
|
||||
re.DOTALL,
|
||||
)
|
||||
else:
|
||||
pattern = re.compile(
|
||||
re.escape(start) + r".*?" + re.escape(end),
|
||||
re.DOTALL,
|
||||
)
|
||||
|
||||
matches = pattern.findall(text)
|
||||
if not matches:
|
||||
return [], text
|
||||
|
||||
remaining = pattern.sub(" ", text).strip()
|
||||
calls = []
|
||||
for match in matches:
|
||||
call_body = match.strip().removeprefix(start)
|
||||
if end:
|
||||
call_body = call_body.removesuffix(end)
|
||||
call_body = call_body.strip()
|
||||
try:
|
||||
parsed = parse_fn(call_body, tools)
|
||||
except Exception as e:
|
||||
print(
|
||||
f"[mlx_utils] Invalid tool call: {call_body!r} ({e})",
|
||||
file=sys.stderr,
|
||||
)
|
||||
continue
|
||||
if not isinstance(parsed, list):
|
||||
parsed = [parsed]
|
||||
for tc in parsed:
|
||||
calls.append(
|
||||
{
|
||||
"index": len(calls),
|
||||
"id": str(uuid.uuid4()),
|
||||
"name": (tc.get("name") or "").strip(),
|
||||
"arguments": json.dumps(tc.get("arguments", {}), ensure_ascii=False),
|
||||
}
|
||||
)
|
||||
return calls, remaining
|
||||
65
backend/python/common/python_utils.py
Normal file
65
backend/python/common/python_utils.py
Normal file
@@ -0,0 +1,65 @@
|
||||
"""Generic utilities shared across Python gRPC backends.
|
||||
|
||||
These helpers don't depend on any specific inference framework and can be
|
||||
imported by any backend that needs to parse LocalAI gRPC options or build a
|
||||
chat-template-compatible message list from proto Message objects.
|
||||
"""
|
||||
import json
|
||||
|
||||
|
||||
def parse_options(options_list):
|
||||
"""Parse Options[] list of ``key:value`` strings into a dict.
|
||||
|
||||
Supports type inference for common cases (bool, int, float). Unknown or
|
||||
mixed-case values are returned as strings.
|
||||
|
||||
Used by LoadModel to extract backend-specific options passed via
|
||||
``ModelOptions.Options`` in ``backend.proto``.
|
||||
"""
|
||||
opts = {}
|
||||
for opt in options_list:
|
||||
if ":" not in opt:
|
||||
continue
|
||||
key, value = opt.split(":", 1)
|
||||
key = key.strip()
|
||||
value = value.strip()
|
||||
# Try type conversion
|
||||
if value.lower() in ("true", "false"):
|
||||
opts[key] = value.lower() == "true"
|
||||
else:
|
||||
try:
|
||||
opts[key] = int(value)
|
||||
except ValueError:
|
||||
try:
|
||||
opts[key] = float(value)
|
||||
except ValueError:
|
||||
opts[key] = value
|
||||
return opts
|
||||
|
||||
|
||||
def messages_to_dicts(proto_messages):
|
||||
"""Convert proto ``Message`` objects to dicts suitable for ``apply_chat_template``.
|
||||
|
||||
Handles: ``role``, ``content``, ``name``, ``tool_call_id``,
|
||||
``reasoning_content``, ``tool_calls`` (JSON string → Python list).
|
||||
|
||||
HuggingFace chat templates (and their MLX/vLLM wrappers) expect a list of
|
||||
plain dicts — proto Message objects don't work directly with Jinja, so
|
||||
this conversion is needed before every ``apply_chat_template`` call.
|
||||
"""
|
||||
result = []
|
||||
for msg in proto_messages:
|
||||
d = {"role": msg.role, "content": msg.content or ""}
|
||||
if msg.name:
|
||||
d["name"] = msg.name
|
||||
if msg.tool_call_id:
|
||||
d["tool_call_id"] = msg.tool_call_id
|
||||
if msg.reasoning_content:
|
||||
d["reasoning_content"] = msg.reasoning_content
|
||||
if msg.tool_calls:
|
||||
try:
|
||||
d["tool_calls"] = json.loads(msg.tool_calls)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
result.append(d)
|
||||
return result
|
||||
@@ -1,63 +1,22 @@
|
||||
"""Shared utilities for vLLM-based backends."""
|
||||
import json
|
||||
"""vLLM-specific helpers for the vllm and vllm-omni gRPC backends.
|
||||
|
||||
Generic helpers (``parse_options``, ``messages_to_dicts``) live in
|
||||
``python_utils`` and are re-exported here for backwards compatibility with
|
||||
existing imports in both backends.
|
||||
"""
|
||||
import sys
|
||||
|
||||
from python_utils import messages_to_dicts, parse_options
|
||||
|
||||
def parse_options(options_list):
|
||||
"""Parse Options[] list of 'key:value' strings into a dict.
|
||||
|
||||
Supports type inference for common cases (bool, int, float).
|
||||
Used by LoadModel to extract backend-specific options.
|
||||
"""
|
||||
opts = {}
|
||||
for opt in options_list:
|
||||
if ":" not in opt:
|
||||
continue
|
||||
key, value = opt.split(":", 1)
|
||||
key = key.strip()
|
||||
value = value.strip()
|
||||
# Try type conversion
|
||||
if value.lower() in ("true", "false"):
|
||||
opts[key] = value.lower() == "true"
|
||||
else:
|
||||
try:
|
||||
opts[key] = int(value)
|
||||
except ValueError:
|
||||
try:
|
||||
opts[key] = float(value)
|
||||
except ValueError:
|
||||
opts[key] = value
|
||||
return opts
|
||||
|
||||
|
||||
def messages_to_dicts(proto_messages):
|
||||
"""Convert proto Message objects to list of dicts for apply_chat_template().
|
||||
|
||||
Handles: role, content, name, tool_call_id, reasoning_content, tool_calls (JSON string -> list).
|
||||
"""
|
||||
result = []
|
||||
for msg in proto_messages:
|
||||
d = {"role": msg.role, "content": msg.content or ""}
|
||||
if msg.name:
|
||||
d["name"] = msg.name
|
||||
if msg.tool_call_id:
|
||||
d["tool_call_id"] = msg.tool_call_id
|
||||
if msg.reasoning_content:
|
||||
d["reasoning_content"] = msg.reasoning_content
|
||||
if msg.tool_calls:
|
||||
try:
|
||||
d["tool_calls"] = json.loads(msg.tool_calls)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
result.append(d)
|
||||
return result
|
||||
__all__ = ["parse_options", "messages_to_dicts", "setup_parsers"]
|
||||
|
||||
|
||||
def setup_parsers(opts):
|
||||
"""Return (tool_parser_cls, reasoning_parser_cls) tuple from opts dict.
|
||||
"""Return ``(tool_parser_cls, reasoning_parser_cls)`` from an opts dict.
|
||||
|
||||
Uses vLLM's native ToolParserManager and ReasoningParserManager.
|
||||
Returns (None, None) if vLLM is not installed or parsers not available.
|
||||
Uses vLLM's native ``ToolParserManager`` / ``ReasoningParserManager``.
|
||||
Returns ``(None, None)`` if vLLM isn't installed or the requested
|
||||
parser name can't be resolved.
|
||||
"""
|
||||
tool_parser_cls = None
|
||||
reasoning_parser_cls = None
|
||||
|
||||
16
backend/python/insightface/Makefile
Normal file
16
backend/python/insightface/Makefile
Normal file
@@ -0,0 +1,16 @@
|
||||
.DEFAULT_GOAL := install
|
||||
|
||||
.PHONY: install
|
||||
install:
|
||||
bash install.sh
|
||||
|
||||
.PHONY: protogen-clean
|
||||
protogen-clean:
|
||||
$(RM) backend_pb2_grpc.py backend_pb2.py
|
||||
|
||||
.PHONY: clean
|
||||
clean: protogen-clean
|
||||
rm -rf venv __pycache__
|
||||
|
||||
test: install
|
||||
bash test.sh
|
||||
67
backend/python/insightface/README.md
Normal file
67
backend/python/insightface/README.md
Normal file
@@ -0,0 +1,67 @@
|
||||
# insightface backend (LocalAI)
|
||||
|
||||
Face recognition backend backed by ONNX Runtime. Provides face
|
||||
verification (1:1), face analysis (age/gender), face detection, face
|
||||
embedding, and — via LocalAI's built-in vector store — 1:N
|
||||
identification.
|
||||
|
||||
## Engines
|
||||
|
||||
This backend ships with **two** interchangeable engines selected via
|
||||
`LoadModel.Options["engine"]`:
|
||||
|
||||
| engine | Implementation | Models | License |
|
||||
|---|---|---|---|
|
||||
| `insightface` (default) | `insightface.app.FaceAnalysis` | `buffalo_l`, `buffalo_s`, `antelopev2` | **Non-commercial research use only** |
|
||||
| `onnx_direct` | OpenCV `FaceDetectorYN` + `FaceRecognizerSF` | OpenCV Zoo YuNet + SFace | Apache 2.0 (commercial-safe) |
|
||||
|
||||
Both engines implement the same `FaceEngine` protocol in `engines.py`,
|
||||
so the gRPC servicer in `backend.py` doesn't need to know which one is
|
||||
active.
|
||||
|
||||
## LoadModel options
|
||||
|
||||
Common:
|
||||
|
||||
| option | default | description |
|
||||
|---|---|---|
|
||||
| `engine` | `insightface` | one of `insightface`, `onnx_direct` |
|
||||
| `det_size` | `640x640` (insightface), `320x320` (onnx_direct) | detector input size |
|
||||
| `det_thresh` | `0.5` | detector confidence threshold |
|
||||
| `verify_threshold` | `0.35` | default cosine distance cutoff for FaceVerify |
|
||||
|
||||
`insightface` engine:
|
||||
|
||||
| option | default | description |
|
||||
|---|---|---|
|
||||
| `model_pack` | `buffalo_l` | which insightface pack to load |
|
||||
|
||||
`onnx_direct` engine:
|
||||
|
||||
| option | default | description |
|
||||
|---|---|---|
|
||||
| `detector_onnx` | *(required)* | path to YuNet-compatible ONNX |
|
||||
| `recognizer_onnx` | *(required)* | path to SFace-compatible ONNX |
|
||||
|
||||
## Adding a new model pack
|
||||
|
||||
1. If it's an insightface pack (auto-downloadable or manually extracted
|
||||
into `~/.insightface/models/<name>/`), just add a new gallery entry
|
||||
in `backend/index.yaml` with `options: ["engine:insightface",
|
||||
"model_pack:<name>"]`. No code change.
|
||||
2. If it's an Apache-licensed ONNX pair, add a gallery entry with
|
||||
`options: ["engine:onnx_direct", "detector_onnx:...",
|
||||
"recognizer_onnx:..."]`. If the detector or recognizer has a
|
||||
different input-tensor shape than YuNet/SFace, you may need a new
|
||||
engine implementation in `engines.py`; the two-engine seam makes
|
||||
that a self-contained change.
|
||||
|
||||
## Running tests locally
|
||||
|
||||
```bash
|
||||
make -C backend/python/insightface # install deps + bake models
|
||||
make -C backend/python/insightface test # run test.py
|
||||
```
|
||||
|
||||
The OpenCV Zoo tests skip gracefully when `/models/opencv/*.onnx` is
|
||||
absent (e.g. on dev boxes where `install.sh` wasn't run).
|
||||
312
backend/python/insightface/backend.py
Normal file
312
backend/python/insightface/backend.py
Normal file
@@ -0,0 +1,312 @@
|
||||
#!/usr/bin/env python3
|
||||
"""gRPC server for the insightface face recognition backend.
|
||||
|
||||
Implements Health / LoadModel / Status plus the face-specific methods:
|
||||
Embedding, Detect, FaceVerify, FaceAnalyze. The heavy lifting is
|
||||
delegated to engines.py — this file is just the gRPC plumbing.
|
||||
"""
|
||||
import argparse
|
||||
import base64
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
import time
|
||||
from concurrent import futures
|
||||
from io import BytesIO
|
||||
|
||||
import backend_pb2
|
||||
import backend_pb2_grpc
|
||||
import cv2
|
||||
import grpc
|
||||
import numpy as np
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "common"))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "common"))
|
||||
from grpc_auth import get_auth_interceptors # noqa: E402
|
||||
|
||||
from engines import FaceEngine, build_engine # noqa: E402
|
||||
|
||||
_ONE_DAY = 60 * 60 * 24
|
||||
MAX_WORKERS = int(os.environ.get("PYTHON_GRPC_MAX_WORKERS", "1"))
|
||||
|
||||
# Default cosine-distance threshold for "same person" on buffalo_l
|
||||
# ArcFace R50. Clients can override per-request; clients using SFace
|
||||
# should pass threshold≈0.4 since the distance distribution is wider.
|
||||
DEFAULT_VERIFY_THRESHOLD = 0.35
|
||||
|
||||
|
||||
def _decode_image(src: str) -> np.ndarray | None:
|
||||
"""Decode a base64-encoded image into an OpenCV BGR numpy array."""
|
||||
if not src:
|
||||
return None
|
||||
try:
|
||||
data = base64.b64decode(src, validate=False)
|
||||
except Exception:
|
||||
return None
|
||||
arr = np.frombuffer(data, dtype=np.uint8)
|
||||
if arr.size == 0:
|
||||
return None
|
||||
img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
|
||||
return img
|
||||
|
||||
|
||||
def _parse_options(raw: list[str]) -> dict[str, str]:
|
||||
out: dict[str, str] = {}
|
||||
for entry in raw:
|
||||
if ":" not in entry:
|
||||
continue
|
||||
k, v = entry.split(":", 1)
|
||||
out[k.strip()] = v.strip()
|
||||
return out
|
||||
|
||||
|
||||
class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
def __init__(self) -> None:
|
||||
self.engine: FaceEngine | None = None
|
||||
self.engine_name: str = ""
|
||||
self.model_name: str = ""
|
||||
self.verify_threshold: float = DEFAULT_VERIFY_THRESHOLD
|
||||
|
||||
def Health(self, request, context):
|
||||
return backend_pb2.Reply(message=bytes("OK", "utf-8"))
|
||||
|
||||
def LoadModel(self, request, context):
|
||||
options = _parse_options(list(request.Options))
|
||||
# Surface LocalAI's models directory (ModelPath) so engines can
|
||||
# anchor relative paths — OnnxDirectEngine's detector_onnx /
|
||||
# recognizer_onnx point at gallery-managed files that LocalAI
|
||||
# dropped there, and InsightFaceEngine auto-downloads its packs
|
||||
# into that same directory alongside every other managed model.
|
||||
# Private key to avoid clashing with user-provided options.
|
||||
if request.ModelPath:
|
||||
options["_model_dir"] = request.ModelPath
|
||||
|
||||
engine_name = options.get("engine", "insightface")
|
||||
try:
|
||||
self.engine = build_engine(engine_name)
|
||||
self.engine.prepare(options)
|
||||
except Exception as err: # pragma: no cover - exercised via e2e
|
||||
return backend_pb2.Result(success=False, message=f"Failed to load face engine: {err}")
|
||||
|
||||
self.engine_name = engine_name
|
||||
self.model_name = request.Model or options.get("model_pack", "")
|
||||
if "verify_threshold" in options:
|
||||
try:
|
||||
self.verify_threshold = float(options["verify_threshold"])
|
||||
except ValueError:
|
||||
pass
|
||||
print(f"[insightface] engine={engine_name} model={self.model_name} loaded", file=sys.stderr)
|
||||
return backend_pb2.Result(success=True, message="Model loaded successfully")
|
||||
|
||||
def Status(self, request, context):
|
||||
state = (
|
||||
backend_pb2.StatusResponse.READY
|
||||
if self.engine is not None
|
||||
else backend_pb2.StatusResponse.UNINITIALIZED
|
||||
)
|
||||
return backend_pb2.StatusResponse(state=state)
|
||||
|
||||
def Embedding(self, request, context):
|
||||
if self.engine is None:
|
||||
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
|
||||
context.set_details("face model not loaded")
|
||||
return backend_pb2.EmbeddingResult()
|
||||
if not request.Images:
|
||||
context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
|
||||
context.set_details("Embedding requires Images[0] to be a base64 image")
|
||||
return backend_pb2.EmbeddingResult()
|
||||
|
||||
img = _decode_image(request.Images[0])
|
||||
if img is None:
|
||||
context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
|
||||
context.set_details("failed to decode image")
|
||||
return backend_pb2.EmbeddingResult()
|
||||
|
||||
vec = self.engine.embed(img)
|
||||
if vec is None:
|
||||
context.set_code(grpc.StatusCode.NOT_FOUND)
|
||||
context.set_details("no face detected")
|
||||
return backend_pb2.EmbeddingResult()
|
||||
return backend_pb2.EmbeddingResult(embeddings=[float(x) for x in vec])
|
||||
|
||||
def Detect(self, request, context):
|
||||
if self.engine is None:
|
||||
return backend_pb2.DetectResponse()
|
||||
img = _decode_image(request.src)
|
||||
if img is None:
|
||||
return backend_pb2.DetectResponse()
|
||||
detections = []
|
||||
for d in self.engine.detect(img):
|
||||
x1, y1, x2, y2 = d.bbox
|
||||
detections.append(
|
||||
backend_pb2.Detection(
|
||||
x=float(x1),
|
||||
y=float(y1),
|
||||
width=float(x2 - x1),
|
||||
height=float(y2 - y1),
|
||||
confidence=float(d.score),
|
||||
class_name="face",
|
||||
)
|
||||
)
|
||||
return backend_pb2.DetectResponse(Detections=detections)
|
||||
|
||||
def FaceVerify(self, request, context):
|
||||
if self.engine is None:
|
||||
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
|
||||
context.set_details("face model not loaded")
|
||||
return backend_pb2.FaceVerifyResponse()
|
||||
|
||||
img1 = _decode_image(request.img1)
|
||||
img2 = _decode_image(request.img2)
|
||||
if img1 is None or img2 is None:
|
||||
context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
|
||||
context.set_details("failed to decode one or both images")
|
||||
return backend_pb2.FaceVerifyResponse()
|
||||
|
||||
threshold = request.threshold if request.threshold > 0 else self.verify_threshold
|
||||
|
||||
start = time.time()
|
||||
e1 = self.engine.embed(img1)
|
||||
e2 = self.engine.embed(img2)
|
||||
if e1 is None or e2 is None:
|
||||
context.set_code(grpc.StatusCode.NOT_FOUND)
|
||||
context.set_details("no face detected in one or both images")
|
||||
return backend_pb2.FaceVerifyResponse()
|
||||
|
||||
# Both engines return L2-normalized vectors, so the dot product
|
||||
# is the cosine similarity directly.
|
||||
sim = float(np.dot(e1, e2))
|
||||
distance = 1.0 - sim
|
||||
verified = distance < threshold
|
||||
confidence = max(0.0, min(100.0, (1.0 - distance / threshold) * 100.0)) if threshold > 0 else 0.0
|
||||
|
||||
# Detect once per image — region is needed for the response and
|
||||
# potentially for the antispoof crop. Returns the highest-score face.
|
||||
def _best_detection(img):
|
||||
dets = self.engine.detect(img)
|
||||
if not dets:
|
||||
return None
|
||||
return max(dets, key=lambda d: d.score)
|
||||
|
||||
def _region(det) -> backend_pb2.FacialArea:
|
||||
if det is None:
|
||||
return backend_pb2.FacialArea()
|
||||
x1, y1, x2, y2 = det.bbox
|
||||
return backend_pb2.FacialArea(x=x1, y=y1, w=x2 - x1, h=y2 - y1)
|
||||
|
||||
det1 = _best_detection(img1)
|
||||
det2 = _best_detection(img2)
|
||||
|
||||
img1_is_real = False
|
||||
img1_score = 0.0
|
||||
img2_is_real = False
|
||||
img2_score = 0.0
|
||||
if request.anti_spoofing:
|
||||
spoof1 = self.engine.antispoof(img1, det1.bbox) if det1 is not None else None
|
||||
spoof2 = self.engine.antispoof(img2, det2.bbox) if det2 is not None else None
|
||||
if spoof1 is None or spoof2 is None:
|
||||
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
|
||||
context.set_details(
|
||||
"anti_spoofing requested but no antispoof model is loaded — "
|
||||
"install `silent-face-antispoofing` or pick a gallery entry "
|
||||
"that bundles MiniFASNet weights"
|
||||
)
|
||||
return backend_pb2.FaceVerifyResponse()
|
||||
img1_is_real, img1_score = spoof1.is_real, spoof1.score
|
||||
img2_is_real, img2_score = spoof2.is_real, spoof2.score
|
||||
# Failed liveness vetoes verification regardless of similarity.
|
||||
if not (img1_is_real and img2_is_real):
|
||||
verified = False
|
||||
|
||||
return backend_pb2.FaceVerifyResponse(
|
||||
verified=verified,
|
||||
distance=float(distance),
|
||||
threshold=float(threshold),
|
||||
confidence=float(confidence),
|
||||
model=self.model_name or self.engine_name,
|
||||
img1_area=_region(det1),
|
||||
img2_area=_region(det2),
|
||||
processing_time_ms=float((time.time() - start) * 1000.0),
|
||||
img1_is_real=img1_is_real,
|
||||
img1_antispoof_score=float(img1_score),
|
||||
img2_is_real=img2_is_real,
|
||||
img2_antispoof_score=float(img2_score),
|
||||
)
|
||||
|
||||
def FaceAnalyze(self, request, context):
|
||||
if self.engine is None:
|
||||
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
|
||||
context.set_details("face model not loaded")
|
||||
return backend_pb2.FaceAnalyzeResponse()
|
||||
img = _decode_image(request.img)
|
||||
if img is None:
|
||||
context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
|
||||
context.set_details("failed to decode image")
|
||||
return backend_pb2.FaceAnalyzeResponse()
|
||||
|
||||
faces = []
|
||||
for attrs in self.engine.analyze(img):
|
||||
x, y, w, h = attrs.region
|
||||
fa = backend_pb2.FaceAnalysis(
|
||||
region=backend_pb2.FacialArea(x=float(x), y=float(y), w=float(w), h=float(h)),
|
||||
face_confidence=float(attrs.face_confidence),
|
||||
)
|
||||
if attrs.age is not None:
|
||||
fa.age = float(attrs.age)
|
||||
if attrs.dominant_gender:
|
||||
fa.dominant_gender = attrs.dominant_gender
|
||||
for k, v in attrs.gender.items():
|
||||
fa.gender[k] = float(v)
|
||||
if request.anti_spoofing:
|
||||
bbox = (float(x), float(y), float(x + w), float(y + h))
|
||||
spoof = self.engine.antispoof(img, bbox)
|
||||
if spoof is None:
|
||||
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
|
||||
context.set_details(
|
||||
"anti_spoofing requested but no antispoof model is loaded — "
|
||||
"install `silent-face-antispoofing` or pick a gallery entry "
|
||||
"that bundles MiniFASNet weights"
|
||||
)
|
||||
return backend_pb2.FaceAnalyzeResponse()
|
||||
fa.is_real = spoof.is_real
|
||||
fa.antispoof_score = float(spoof.score)
|
||||
faces.append(fa)
|
||||
return backend_pb2.FaceAnalyzeResponse(faces=faces)
|
||||
|
||||
|
||||
def serve(address: str) -> None:
|
||||
server = grpc.server(
|
||||
futures.ThreadPoolExecutor(max_workers=MAX_WORKERS),
|
||||
options=[
|
||||
("grpc.max_message_length", 50 * 1024 * 1024),
|
||||
("grpc.max_send_message_length", 50 * 1024 * 1024),
|
||||
("grpc.max_receive_message_length", 50 * 1024 * 1024),
|
||||
],
|
||||
interceptors=get_auth_interceptors(),
|
||||
)
|
||||
backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
|
||||
server.add_insecure_port(address)
|
||||
server.start()
|
||||
print("[insightface] Server started. Listening on: " + address, file=sys.stderr)
|
||||
|
||||
def _stop(sig, frame): # pragma: no cover
|
||||
print("[insightface] shutting down")
|
||||
server.stop(0)
|
||||
sys.exit(0)
|
||||
|
||||
signal.signal(signal.SIGINT, _stop)
|
||||
signal.signal(signal.SIGTERM, _stop)
|
||||
|
||||
try:
|
||||
while True:
|
||||
time.sleep(_ONE_DAY)
|
||||
except KeyboardInterrupt:
|
||||
server.stop(0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Run the insightface gRPC server.")
|
||||
parser.add_argument("--addr", default="localhost:50051", help="The address to bind the server to.")
|
||||
args = parser.parse_args()
|
||||
print(f"[insightface] startup: {args}", file=sys.stderr)
|
||||
serve(args.addr)
|
||||
573
backend/python/insightface/engines.py
Normal file
573
backend/python/insightface/engines.py
Normal file
@@ -0,0 +1,573 @@
|
||||
"""Face recognition engine implementations for the LocalAI insightface backend.
|
||||
|
||||
Two engines are provided:
|
||||
|
||||
* InsightFaceEngine — wraps insightface.app.FaceAnalysis. Supports
|
||||
buffalo_l / buffalo_s / antelopev2 model packs
|
||||
with SCRFD detector + ArcFace recognizer +
|
||||
genderage head. NON-COMMERCIAL research use
|
||||
only (upstream license).
|
||||
|
||||
* OnnxDirectEngine — loads detector + recognizer ONNX files directly
|
||||
via onnxruntime. Used for OpenCV Zoo models
|
||||
(YuNet + SFace) and any future Apache-licensed
|
||||
model set. Does not support analyze().
|
||||
|
||||
Both engines expose the same interface so the gRPC servicer (backend.py)
|
||||
can dispatch without knowing which one is active.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Protocol
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
@dataclass
|
||||
class FaceDetection:
|
||||
bbox: tuple[float, float, float, float] # x1, y1, x2, y2
|
||||
score: float
|
||||
landmarks: np.ndarray | None = None # 5x2 keypoints when available
|
||||
|
||||
|
||||
@dataclass
|
||||
class FaceAttributes:
|
||||
region: tuple[float, float, float, float] # x, y, w, h
|
||||
face_confidence: float
|
||||
age: float | None = None
|
||||
dominant_gender: str | None = None
|
||||
gender: dict[str, float] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SpoofResult:
|
||||
is_real: bool
|
||||
score: float # averaged probability of the "real" class, 0.0-1.0
|
||||
|
||||
|
||||
class FaceEngine(Protocol):
|
||||
"""Minimal interface every engine must implement."""
|
||||
|
||||
def prepare(self, options: dict[str, str]) -> None: ...
|
||||
def detect(self, img: np.ndarray) -> list[FaceDetection]: ...
|
||||
def embed(self, img: np.ndarray) -> np.ndarray | None: ...
|
||||
def analyze(self, img: np.ndarray) -> list[FaceAttributes]: ...
|
||||
# Optional: returns None when no antispoof model is loaded.
|
||||
def antispoof(self, img: np.ndarray, bbox: tuple[float, float, float, float]) -> SpoofResult | None: ...
|
||||
|
||||
|
||||
# ─── Antispoofer (Silent-Face MiniFASNet) ──────────────────────────────
|
||||
|
||||
class Antispoofer:
|
||||
"""Liveness detector using the Silent-Face MiniFASNet ensemble.
|
||||
|
||||
Loads up to two ONNX exports (MiniFASNetV2 at scale 2.7 and
|
||||
MiniFASNetV1SE at scale 4.0). Both are 80x80 BGR-float32-input
|
||||
classifiers with 3 output logits where index 1 = "real". When both
|
||||
are loaded, softmax outputs are averaged before argmax — the same
|
||||
ensembling the upstream `test.py` does.
|
||||
|
||||
Preprocessing matches yakhyo/face-anti-spoofing's reference impl:
|
||||
each model gets its own scale-expanded crop centered on the face
|
||||
bbox, resized to 80x80, fed straight as float32 BGR (no /255, no
|
||||
mean/std). See `_crop_face` for the bbox math.
|
||||
|
||||
A single model also works (the missing one is simply skipped).
|
||||
"""
|
||||
|
||||
INPUT_SIZE = (80, 80) # h, w
|
||||
REAL_CLASS_IDX = 1
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._sessions: list[tuple[Any, float, str, str]] = [] # (session, scale, input_name, output_name)
|
||||
self.threshold: float = 0.5
|
||||
|
||||
def load(self, model_paths: list[tuple[str, float]], threshold: float = 0.5) -> None:
|
||||
"""Load one or more (path, scale) pairs."""
|
||||
import onnxruntime as ort
|
||||
|
||||
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||||
for path, scale in model_paths:
|
||||
session = ort.InferenceSession(path, providers=providers)
|
||||
input_name = session.get_inputs()[0].name
|
||||
output_name = session.get_outputs()[0].name
|
||||
self._sessions.append((session, float(scale), input_name, output_name))
|
||||
self.threshold = float(threshold)
|
||||
|
||||
@property
|
||||
def loaded(self) -> bool:
|
||||
return bool(self._sessions)
|
||||
|
||||
def _crop_face(self, img: np.ndarray, bbox: tuple[float, float, float, float], scale: float) -> np.ndarray:
|
||||
# bbox is (x1, y1, x2, y2) in source-image coordinates.
|
||||
src_h, src_w = img.shape[:2]
|
||||
x1, y1, x2, y2 = bbox
|
||||
box_w = max(1.0, x2 - x1)
|
||||
box_h = max(1.0, y2 - y1)
|
||||
|
||||
# Clamp scale so the expanded crop fits inside the source image.
|
||||
scale = min((src_h - 1) / box_h, (src_w - 1) / box_w, scale)
|
||||
new_w = box_w * scale
|
||||
new_h = box_h * scale
|
||||
|
||||
cx = x1 + box_w / 2.0
|
||||
cy = y1 + box_h / 2.0
|
||||
|
||||
cx1 = max(0, int(cx - new_w / 2.0))
|
||||
cy1 = max(0, int(cy - new_h / 2.0))
|
||||
cx2 = min(src_w - 1, int(cx + new_w / 2.0))
|
||||
cy2 = min(src_h - 1, int(cy + new_h / 2.0))
|
||||
|
||||
cropped = img[cy1 : cy2 + 1, cx1 : cx2 + 1]
|
||||
if cropped.size == 0:
|
||||
cropped = img
|
||||
out_h, out_w = self.INPUT_SIZE
|
||||
return cv2.resize(cropped, (out_w, out_h))
|
||||
|
||||
@staticmethod
|
||||
def _softmax(x: np.ndarray) -> np.ndarray:
|
||||
e = np.exp(x - np.max(x, axis=1, keepdims=True))
|
||||
return e / e.sum(axis=1, keepdims=True)
|
||||
|
||||
def predict(self, img: np.ndarray, bbox: tuple[float, float, float, float]) -> SpoofResult:
|
||||
if not self._sessions:
|
||||
raise RuntimeError("Antispoofer.predict called with no models loaded")
|
||||
accum = np.zeros((1, 3), dtype=np.float32)
|
||||
for session, scale, input_name, output_name in self._sessions:
|
||||
face = self._crop_face(img, bbox, scale).astype(np.float32)
|
||||
tensor = np.transpose(face, (2, 0, 1))[np.newaxis, ...]
|
||||
logits = session.run([output_name], {input_name: tensor})[0]
|
||||
accum += self._softmax(logits)
|
||||
accum /= float(len(self._sessions))
|
||||
real_prob = float(accum[0, self.REAL_CLASS_IDX])
|
||||
is_real = int(np.argmax(accum)) == self.REAL_CLASS_IDX and real_prob >= self.threshold
|
||||
return SpoofResult(is_real=is_real, score=real_prob)
|
||||
|
||||
|
||||
def _build_antispoofer(options: dict[str, str], model_dir: str | None) -> Antispoofer | None:
|
||||
"""Instantiate an Antispoofer from option keys, or return None.
|
||||
|
||||
Recognised options:
|
||||
antispoof_v2_onnx — path/filename of MiniFASNetV2 (scale 2.7)
|
||||
antispoof_v1se_onnx — path/filename of MiniFASNetV1SE (scale 4.0)
|
||||
antispoof_threshold — real-class probability threshold, default 0.5
|
||||
|
||||
Either or both can be provided. Returns None when neither is set.
|
||||
"""
|
||||
pairs: list[tuple[str, float]] = []
|
||||
v2 = options.get("antispoof_v2_onnx", "")
|
||||
if v2:
|
||||
pairs.append((_resolve_model_path(v2, model_dir=model_dir), 2.7))
|
||||
v1se = options.get("antispoof_v1se_onnx", "")
|
||||
if v1se:
|
||||
pairs.append((_resolve_model_path(v1se, model_dir=model_dir), 4.0))
|
||||
if not pairs:
|
||||
return None
|
||||
threshold = float(options.get("antispoof_threshold", "0.5"))
|
||||
spoofer = Antispoofer()
|
||||
spoofer.load(pairs, threshold=threshold)
|
||||
return spoofer
|
||||
|
||||
|
||||
# ─── InsightFaceEngine ────────────────────────────────────────────────
|
||||
|
||||
# Canonical ONNX manifest for each upstream insightface pack (v0.7 release
|
||||
# at github.com/deepinsight/insightface/releases). LocalAI's gallery extracts
|
||||
# these zips flat into the models directory, so when multiple packs or other
|
||||
# backends drop their own ONNX files alongside, the glob-the-directory
|
||||
# approach picks up foreign files and insightface's model_zoo.get_model()
|
||||
# raises IndexError trying to index `input_shape[2]` on a tensor that isn't
|
||||
# shaped like a face model. The manifest lets us pre-filter to only the
|
||||
# files that actually belong to the requested pack — deterministic, correct
|
||||
# pack choice, no crashes on neighbour ONNX files.
|
||||
_KNOWN_PACK_MANIFESTS: dict[str, frozenset[str]] = {
|
||||
"buffalo_l": frozenset({
|
||||
"det_10g.onnx",
|
||||
"w600k_r50.onnx",
|
||||
"genderage.onnx",
|
||||
"2d106det.onnx",
|
||||
"1k3d68.onnx",
|
||||
}),
|
||||
"buffalo_sc": frozenset({
|
||||
"det_500m.onnx",
|
||||
"w600k_mbf.onnx",
|
||||
}),
|
||||
}
|
||||
|
||||
|
||||
class InsightFaceEngine:
|
||||
"""Drives insightface's model_zoo directly — no FaceAnalysis wrapper.
|
||||
|
||||
FaceAnalysis is a thin 50-line orchestration (glob for ONNX files
|
||||
in `<root>/models/<name>/`, route each through `model_zoo.get_model`,
|
||||
build a `{taskname: model}` dict, then loop per-face at inference).
|
||||
We reimplement the same loop here so we can:
|
||||
|
||||
1. Load packs from whatever directory LocalAI's gallery extracted
|
||||
them into — flat (buffalo_l/s/sc — ONNX at `<dir>/*.onnx`) or
|
||||
nested (buffalo_m/antelopev2 — ONNX at `<dir>/<name>/*.onnx`)
|
||||
without needing a specific layout on disk.
|
||||
2. Skip insightface's built-in auto-download entirely: weight
|
||||
delivery is LocalAI's gallery `files:` job now, checksum-
|
||||
verified and cached alongside every other managed model.
|
||||
|
||||
The actual inference classes (RetinaFace, ArcFaceONNX, Attribute,
|
||||
Landmark) stay in insightface — we only reimplement the ~50 lines
|
||||
of glue around them.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.models: dict[str, Any] = {}
|
||||
self.det_model: Any = None
|
||||
self.model_pack: str = "buffalo_l"
|
||||
self.det_size: tuple[int, int] = (640, 640)
|
||||
self.det_thresh: float = 0.5
|
||||
self._providers: list[str] = ["CPUExecutionProvider"]
|
||||
self._antispoofer: Antispoofer | None = None
|
||||
|
||||
def prepare(self, options: dict[str, str]) -> None:
|
||||
import glob
|
||||
import os
|
||||
|
||||
from insightface.model_zoo import model_zoo
|
||||
|
||||
self.model_pack = options.get("model_pack", "buffalo_l")
|
||||
self.det_size = _parse_det_size(options.get("det_size", "640x640"))
|
||||
self.det_thresh = float(options.get("det_thresh", "0.5"))
|
||||
self._antispoofer = _build_antispoofer(options, options.get("_model_dir"))
|
||||
|
||||
pack_dir = _locate_insightface_pack(options, self.model_pack)
|
||||
if pack_dir is None:
|
||||
raise ValueError(
|
||||
f"no insightface pack '{self.model_pack}' found — install via "
|
||||
f"`local-ai models install insightface-{self.model_pack.replace('_', '-')}`"
|
||||
)
|
||||
|
||||
onnx_files = sorted(glob.glob(os.path.join(pack_dir, "*.onnx")))
|
||||
# When the pack extracts flat into a shared models directory it
|
||||
# mixes with ONNX files from other backends (opencv face engine,
|
||||
# MiniFASNet antispoof, WeSpeaker voice embedding, other buffalo
|
||||
# packs installed earlier). Feeding those into model_zoo.get_model()
|
||||
# blows up inside insightface's router — it assumes a 4-D NCHW
|
||||
# input and indexes `input_shape[2]` on tensors that aren't shaped
|
||||
# like a face model, raising IndexError. For the upstream packs we
|
||||
# know the exact ONNX manifest; scoping to it makes the load
|
||||
# deterministic (without it, det_10g.onnx from buffalo_l sorts
|
||||
# before det_500m.onnx from buffalo_sc and silently wins).
|
||||
manifest = _KNOWN_PACK_MANIFESTS.get(self.model_pack)
|
||||
if manifest is not None:
|
||||
scoped = [f for f in onnx_files if os.path.basename(f) in manifest]
|
||||
if scoped:
|
||||
onnx_files = scoped
|
||||
if not onnx_files:
|
||||
raise ValueError(f"no ONNX files in pack directory: {pack_dir}")
|
||||
|
||||
# CUDAExecutionProvider is picked automatically by onnxruntime-gpu
|
||||
# when available; falling back to CPU keeps the CPU-only image
|
||||
# working. ctx_id=0 means "first GPU if any, else CPU".
|
||||
self._providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||||
|
||||
self.models = {}
|
||||
skipped: list[tuple[str, str]] = []
|
||||
for onnx_file in onnx_files:
|
||||
try:
|
||||
m = model_zoo.get_model(onnx_file, providers=self._providers)
|
||||
except Exception as err:
|
||||
# Foreign ONNX (wrong rank/shape, non-insightface model) —
|
||||
# older insightface versions raise IndexError / ValueError
|
||||
# instead of returning None. Keep loading the rest.
|
||||
skipped.append((os.path.basename(onnx_file), str(err)))
|
||||
continue
|
||||
if m is None:
|
||||
skipped.append((os.path.basename(onnx_file), "unknown taskname"))
|
||||
continue
|
||||
# First occurrence of each taskname wins (matches FaceAnalysis).
|
||||
if m.taskname not in self.models:
|
||||
self.models[m.taskname] = m
|
||||
|
||||
if skipped:
|
||||
import sys
|
||||
print(
|
||||
f"[insightface] skipped {len(skipped)} non-pack ONNX file(s) in {pack_dir}: "
|
||||
+ ", ".join(f"{n} ({why})" for n, why in skipped),
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
if "detection" not in self.models:
|
||||
raise ValueError(f"no detector (taskname='detection') found in {pack_dir}")
|
||||
self.det_model = self.models["detection"]
|
||||
|
||||
self.det_model.prepare(0, input_size=self.det_size, det_thresh=self.det_thresh)
|
||||
for name, m in self.models.items():
|
||||
if name != "detection":
|
||||
m.prepare(0)
|
||||
|
||||
def _faces(self, img: np.ndarray) -> list[Any]:
|
||||
"""Run detection + all non-detection models per face."""
|
||||
if self.det_model is None:
|
||||
return []
|
||||
from insightface.app.common import Face
|
||||
|
||||
bboxes, kpss = self.det_model.detect(img, max_num=0)
|
||||
if bboxes is None or bboxes.shape[0] == 0:
|
||||
return []
|
||||
faces: list[Any] = []
|
||||
for i in range(bboxes.shape[0]):
|
||||
bbox = bboxes[i, 0:4]
|
||||
det_score = bboxes[i, 4]
|
||||
kps = kpss[i] if kpss is not None else None
|
||||
face = Face(bbox=bbox, kps=kps, det_score=det_score)
|
||||
for name, m in self.models.items():
|
||||
if name == "detection":
|
||||
continue
|
||||
m.get(img, face)
|
||||
faces.append(face)
|
||||
return faces
|
||||
|
||||
def detect(self, img: np.ndarray) -> list[FaceDetection]:
|
||||
return [
|
||||
FaceDetection(
|
||||
bbox=tuple(float(v) for v in f.bbox),
|
||||
score=float(f.det_score),
|
||||
landmarks=np.array(f.kps) if getattr(f, "kps", None) is not None else None,
|
||||
)
|
||||
for f in self._faces(img)
|
||||
]
|
||||
|
||||
def embed(self, img: np.ndarray) -> np.ndarray | None:
|
||||
faces = self._faces(img)
|
||||
if not faces:
|
||||
return None
|
||||
best = max(faces, key=lambda f: float(f.det_score))
|
||||
if getattr(best, "normed_embedding", None) is None:
|
||||
return None
|
||||
return np.asarray(best.normed_embedding, dtype=np.float32)
|
||||
|
||||
def analyze(self, img: np.ndarray) -> list[FaceAttributes]:
|
||||
out: list[FaceAttributes] = []
|
||||
for f in self._faces(img):
|
||||
x1, y1, x2, y2 = (float(v) for v in f.bbox)
|
||||
region = (x1, y1, x2 - x1, y2 - y1)
|
||||
attrs = FaceAttributes(region=region, face_confidence=float(f.det_score))
|
||||
age = getattr(f, "age", None)
|
||||
if age is not None:
|
||||
attrs.age = float(age)
|
||||
gender = getattr(f, "gender", None)
|
||||
if gender is not None:
|
||||
# genderage head emits argmax, not probabilities —
|
||||
# one-hot dict keeps the API stable.
|
||||
attrs.dominant_gender = "Man" if int(gender) == 1 else "Woman"
|
||||
attrs.gender = {
|
||||
"Man": 1.0 if int(gender) == 1 else 0.0,
|
||||
"Woman": 0.0 if int(gender) == 1 else 1.0,
|
||||
}
|
||||
out.append(attrs)
|
||||
return out
|
||||
|
||||
def antispoof(self, img: np.ndarray, bbox: tuple[float, float, float, float]) -> SpoofResult | None:
|
||||
if self._antispoofer is None or not self._antispoofer.loaded:
|
||||
return None
|
||||
return self._antispoofer.predict(img, bbox)
|
||||
|
||||
|
||||
# ─── OnnxDirectEngine ─────────────────────────────────────────────────
|
||||
|
||||
class OnnxDirectEngine:
|
||||
"""Loads detector + recognizer ONNX files directly.
|
||||
|
||||
Supports the OpenCV Zoo YuNet + SFace pair out of the box. YuNet
|
||||
exposes a C++-level API via cv2.FaceDetectorYN which accepts the
|
||||
ONNX file directly; SFace is driven through cv2.FaceRecognizerSF.
|
||||
Both are Apache 2.0 licensed.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.detector_path: str = ""
|
||||
self.recognizer_path: str = ""
|
||||
self.input_size: tuple[int, int] = (320, 320)
|
||||
self.det_thresh: float = 0.5
|
||||
self._detector: Any = None
|
||||
self._recognizer: Any = None
|
||||
self._antispoofer: Antispoofer | None = None
|
||||
|
||||
def prepare(self, options: dict[str, str]) -> None:
|
||||
raw_det = options.get("detector_onnx", "")
|
||||
raw_rec = options.get("recognizer_onnx", "")
|
||||
if not raw_det or not raw_rec:
|
||||
raise ValueError(
|
||||
"onnx_direct engine requires both detector_onnx and recognizer_onnx options"
|
||||
)
|
||||
model_dir = options.get("_model_dir")
|
||||
self.detector_path = _resolve_model_path(raw_det, model_dir=model_dir)
|
||||
self.recognizer_path = _resolve_model_path(raw_rec, model_dir=model_dir)
|
||||
self.input_size = _parse_det_size(options.get("det_size", "320x320"))
|
||||
self.det_thresh = float(options.get("det_thresh", "0.5"))
|
||||
self._antispoofer = _build_antispoofer(options, model_dir)
|
||||
|
||||
# YuNet is a fixed-size detector; size is reset per detect() call to
|
||||
# match the input frame.
|
||||
self._detector = cv2.FaceDetectorYN.create(
|
||||
self.detector_path,
|
||||
"",
|
||||
self.input_size,
|
||||
score_threshold=self.det_thresh,
|
||||
nms_threshold=0.3,
|
||||
top_k=5000,
|
||||
)
|
||||
self._recognizer = cv2.FaceRecognizerSF.create(self.recognizer_path, "")
|
||||
|
||||
def detect(self, img: np.ndarray) -> list[FaceDetection]:
|
||||
if self._detector is None:
|
||||
return []
|
||||
h, w = img.shape[:2]
|
||||
self._detector.setInputSize((w, h))
|
||||
retval, faces = self._detector.detect(img)
|
||||
if faces is None:
|
||||
return []
|
||||
out: list[FaceDetection] = []
|
||||
for row in faces:
|
||||
x, y, fw, fh = float(row[0]), float(row[1]), float(row[2]), float(row[3])
|
||||
# Landmarks at columns 4..13 are (lx1,ly1,...,lx5,ly5).
|
||||
landmarks = np.array(row[4:14], dtype=np.float32).reshape(5, 2) if len(row) >= 14 else None
|
||||
score = float(row[-1])
|
||||
out.append(FaceDetection(bbox=(x, y, x + fw, y + fh), score=score, landmarks=landmarks))
|
||||
return out
|
||||
|
||||
def embed(self, img: np.ndarray) -> np.ndarray | None:
|
||||
if self._detector is None or self._recognizer is None:
|
||||
return None
|
||||
h, w = img.shape[:2]
|
||||
self._detector.setInputSize((w, h))
|
||||
retval, faces = self._detector.detect(img)
|
||||
if faces is None or len(faces) == 0:
|
||||
return None
|
||||
# Pick the highest-score face (last column is score).
|
||||
best = max(faces, key=lambda r: float(r[-1]))
|
||||
aligned = self._recognizer.alignCrop(img, best)
|
||||
feat = self._recognizer.feature(aligned)
|
||||
vec = np.asarray(feat, dtype=np.float32).flatten()
|
||||
# SFace outputs a 128-dim feature; L2-normalize to make dot-product
|
||||
# comparable to buffalo_l's already-normed 512-dim embedding.
|
||||
norm = float(np.linalg.norm(vec))
|
||||
if norm == 0:
|
||||
return None
|
||||
return vec / norm
|
||||
|
||||
def analyze(self, img: np.ndarray) -> list[FaceAttributes]:
|
||||
# OpenCV Zoo does not ship a demographic classifier; report
|
||||
# only the face-detection regions so callers can still see
|
||||
# how many faces were detected.
|
||||
return [
|
||||
FaceAttributes(
|
||||
region=(
|
||||
d.bbox[0],
|
||||
d.bbox[1],
|
||||
d.bbox[2] - d.bbox[0],
|
||||
d.bbox[3] - d.bbox[1],
|
||||
),
|
||||
face_confidence=d.score,
|
||||
)
|
||||
for d in self.detect(img)
|
||||
]
|
||||
|
||||
def antispoof(self, img: np.ndarray, bbox: tuple[float, float, float, float]) -> SpoofResult | None:
|
||||
if self._antispoofer is None or not self._antispoofer.loaded:
|
||||
return None
|
||||
return self._antispoofer.predict(img, bbox)
|
||||
|
||||
|
||||
# ─── helpers ──────────────────────────────────────────────────────────
|
||||
|
||||
def _parse_det_size(raw: str) -> tuple[int, int]:
|
||||
raw = raw.strip().lower().replace(" ", "")
|
||||
if "x" in raw:
|
||||
w, h = raw.split("x", 1)
|
||||
return (int(w), int(h))
|
||||
n = int(raw)
|
||||
return (n, n)
|
||||
|
||||
|
||||
def _locate_insightface_pack(options: dict[str, str], name: str) -> str | None:
|
||||
"""Find the directory holding the insightface pack's ONNX files.
|
||||
|
||||
LocalAI's gallery `files:` extracts the pack zip straight into the
|
||||
models directory. Upstream packs are inconsistent:
|
||||
|
||||
buffalo_l/s/sc — flat zip, ONNX lands at `<models_dir>/*.onnx`
|
||||
buffalo_m, antelopev2 — wrapped zip, ONNX lands at `<models_dir>/<name>/*.onnx`
|
||||
|
||||
We search, in order:
|
||||
1. `<models_dir>/<name>/` — wrapped-zip layout, or insightface's
|
||||
own FaceAnalysis-style `<root>/models/<name>/` layout.
|
||||
2. `<models_dir>/models/<name>/` — insightface's FaceAnalysis
|
||||
auto-download lands here (handy for dev environments that
|
||||
still have old `~/.insightface` caches).
|
||||
3. `<models_dir>/` — flat-zip layout directly in models dir.
|
||||
|
||||
Returns the first directory whose contents include `*.onnx`.
|
||||
"""
|
||||
import glob
|
||||
import os
|
||||
|
||||
model_dir = options.get("_model_dir") or ""
|
||||
explicit_root = options.get("root")
|
||||
|
||||
candidates: list[str] = []
|
||||
if model_dir:
|
||||
candidates.append(os.path.join(model_dir, name))
|
||||
candidates.append(os.path.join(model_dir, "models", name))
|
||||
candidates.append(model_dir)
|
||||
if explicit_root:
|
||||
expanded = os.path.expanduser(explicit_root)
|
||||
candidates.append(os.path.join(expanded, "models", name))
|
||||
candidates.append(os.path.join(expanded, name))
|
||||
candidates.append(expanded)
|
||||
|
||||
for c in candidates:
|
||||
if os.path.isdir(c) and glob.glob(os.path.join(c, "*.onnx")):
|
||||
return c
|
||||
return None
|
||||
|
||||
|
||||
def _resolve_model_path(path: str, model_dir: str | None = None) -> str:
|
||||
"""Resolve an ONNX file path across the paths LocalAI might deliver it from.
|
||||
|
||||
Search order:
|
||||
1. The path itself if it already resolves (absolute, or relative to CWD).
|
||||
2. `model_dir` (typically `os.path.dirname(ModelOptions.ModelFile)`) —
|
||||
this is how LocalAI surfaces gallery-managed files. When the gallery
|
||||
entry lists `files:`, each one lands under the models directory and
|
||||
backends load them via filename anchored by ModelFile.
|
||||
3. `<script_dir>/<path-without-leading-slash>` — covers dev layouts
|
||||
where someone manually dropped weights inside the backend dir.
|
||||
|
||||
If none hit, return the literal input so cv2/insightface surfaces a
|
||||
clearer error naming the actually-attempted path.
|
||||
"""
|
||||
import os
|
||||
|
||||
if os.path.isfile(path):
|
||||
return path
|
||||
stripped = path.lstrip("/")
|
||||
candidates: list[str] = []
|
||||
if model_dir:
|
||||
candidates.append(os.path.join(model_dir, os.path.basename(path)))
|
||||
candidates.append(os.path.join(model_dir, stripped))
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
candidates.append(os.path.join(script_dir, stripped))
|
||||
for c in candidates:
|
||||
if os.path.isfile(c):
|
||||
return c
|
||||
return path
|
||||
|
||||
|
||||
def build_engine(name: str) -> FaceEngine:
|
||||
"""Factory for the engine selected by LoadModel options."""
|
||||
key = name.strip().lower()
|
||||
if key in ("", "insightface"):
|
||||
return InsightFaceEngine()
|
||||
if key in ("onnx_direct", "onnx-direct", "opencv"):
|
||||
return OnnxDirectEngine()
|
||||
raise ValueError(f"unknown engine: {name!r}")
|
||||
28
backend/python/insightface/install.sh
Executable file
28
backend/python/insightface/install.sh
Executable file
@@ -0,0 +1,28 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
backend_dir=$(dirname $0)
|
||||
if [ -d $backend_dir/common ]; then
|
||||
source $backend_dir/common/libbackend.sh
|
||||
else
|
||||
source $backend_dir/../common/libbackend.sh
|
||||
fi
|
||||
|
||||
installRequirements
|
||||
|
||||
# We deliberately do NOT pre-bake any model weights here. Two reasons:
|
||||
#
|
||||
# 1. Weights should follow LocalAI's gallery-managed download flow
|
||||
# like every other backend. For OpenCV Zoo (YuNet + SFace) the
|
||||
# gallery entries in gallery/index.yaml list the ONNX files via
|
||||
# `files:` with URI + SHA-256 — LocalAI fetches them into the
|
||||
# models directory on `local-ai models install`.
|
||||
#
|
||||
# 2. For insightface model packs (buffalo_l, buffalo_s, buffalo_m,
|
||||
# buffalo_sc, antelopev2), upstream distributes zip archives
|
||||
# only (no individual ONNX URLs). We rely on insightface's own
|
||||
# auto-download machinery (`FaceAnalysis(name=<pack>, root=<dir>)`)
|
||||
# at first LoadModel, pointed at a writable directory. This
|
||||
# matches how rfdetr behaves (uses `inference.get_model()`).
|
||||
#
|
||||
# Net effect: the backend image ships only Python deps (~150MB CPU).
|
||||
7
backend/python/insightface/requirements-cpu.txt
Normal file
7
backend/python/insightface/requirements-cpu.txt
Normal file
@@ -0,0 +1,7 @@
|
||||
insightface
|
||||
onnxruntime
|
||||
opencv-python-headless
|
||||
numpy
|
||||
onnx
|
||||
cython
|
||||
scikit-image
|
||||
7
backend/python/insightface/requirements-cublas12.txt
Normal file
7
backend/python/insightface/requirements-cublas12.txt
Normal file
@@ -0,0 +1,7 @@
|
||||
insightface
|
||||
onnxruntime-gpu
|
||||
opencv-python-headless
|
||||
numpy
|
||||
onnx
|
||||
cython
|
||||
scikit-image
|
||||
3
backend/python/insightface/requirements.txt
Normal file
3
backend/python/insightface/requirements.txt
Normal file
@@ -0,0 +1,3 @@
|
||||
grpcio==1.71.0
|
||||
protobuf
|
||||
grpcio-tools
|
||||
9
backend/python/insightface/run.sh
Executable file
9
backend/python/insightface/run.sh
Executable file
@@ -0,0 +1,9 @@
|
||||
#!/bin/bash
|
||||
backend_dir=$(dirname $0)
|
||||
if [ -d $backend_dir/common ]; then
|
||||
source $backend_dir/common/libbackend.sh
|
||||
else
|
||||
source $backend_dir/../common/libbackend.sh
|
||||
fi
|
||||
|
||||
startBackend $@
|
||||
264
backend/python/insightface/smoke.py
Normal file
264
backend/python/insightface/smoke.py
Normal file
@@ -0,0 +1,264 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Smoke-test every face recognition model configuration shipped in the
|
||||
gallery. Simulates what LocalAI does at runtime: for each config, sets
|
||||
up a models directory, fetches any required files via URL (as the
|
||||
gallery's `files:` list would), then loads + detects + embeds via the
|
||||
in-process BackendServicer — matching the gRPC surface end users hit.
|
||||
|
||||
Run inside the built backend image (venv already has insightface /
|
||||
onnxruntime / opencv-python-headless):
|
||||
|
||||
python smoke.py
|
||||
|
||||
Network is required for the insightface packs (fetched via upstream's
|
||||
FaceAnalysis auto-download at first LoadModel) and for downloading
|
||||
the OpenCV Zoo ONNX files on first run.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
import urllib.request
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
sys.path.insert(0, os.path.dirname(__file__))
|
||||
|
||||
import backend_pb2 # noqa: E402
|
||||
from backend import BackendServicer # noqa: E402
|
||||
|
||||
|
||||
# Gallery `files:` for the OpenCV variants — same URIs + SHA-256s as
|
||||
# gallery/index.yaml lists. Tuples: (filename, uri, sha256).
|
||||
OPENCV_FILES = {
|
||||
"fp32": [
|
||||
(
|
||||
"face_detection_yunet_2023mar.onnx",
|
||||
"https://github.com/opencv/opencv_zoo/raw/main/models/face_detection_yunet/face_detection_yunet_2023mar.onnx",
|
||||
"8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4",
|
||||
),
|
||||
(
|
||||
"face_recognition_sface_2021dec.onnx",
|
||||
"https://github.com/opencv/opencv_zoo/raw/main/models/face_recognition_sface/face_recognition_sface_2021dec.onnx",
|
||||
"0ba9fbfa01b5270c96627c4ef784da859931e02f04419c829e83484087c34e79",
|
||||
),
|
||||
],
|
||||
"int8": [
|
||||
(
|
||||
"face_detection_yunet_2023mar_int8.onnx",
|
||||
"https://github.com/opencv/opencv_zoo/raw/main/models/face_detection_yunet/face_detection_yunet_2023mar_int8.onnx",
|
||||
"321aa5a6afabf7ecc46a3d06bfab2b579dc96eb5c3be7edd365fa04502ad9294",
|
||||
),
|
||||
(
|
||||
"face_recognition_sface_2021dec_int8.onnx",
|
||||
"https://github.com/opencv/opencv_zoo/raw/main/models/face_recognition_sface/face_recognition_sface_2021dec_int8.onnx",
|
||||
"2b0e941e6f16cc048c20aee0c8e31f569118f65d702914540f7bfdc14048d78a",
|
||||
),
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
CONFIGS = [
|
||||
{
|
||||
"name": "insightface-buffalo-l",
|
||||
"options": ["engine:insightface", "model_pack:buffalo_l"],
|
||||
"has_analyze": True,
|
||||
"needs_opencv_files": None,
|
||||
},
|
||||
{
|
||||
"name": "insightface-buffalo-sc",
|
||||
"options": ["engine:insightface", "model_pack:buffalo_sc"],
|
||||
# buffalo_sc has recognizer only — no landmarks, no genderage.
|
||||
"has_analyze": False,
|
||||
"needs_opencv_files": None,
|
||||
},
|
||||
{
|
||||
"name": "insightface-buffalo-s",
|
||||
"options": ["engine:insightface", "model_pack:buffalo_s"],
|
||||
"has_analyze": True,
|
||||
"needs_opencv_files": None,
|
||||
},
|
||||
{
|
||||
"name": "insightface-buffalo-m",
|
||||
"options": ["engine:insightface", "model_pack:buffalo_m"],
|
||||
"has_analyze": True,
|
||||
"needs_opencv_files": None,
|
||||
},
|
||||
{
|
||||
"name": "insightface-antelopev2",
|
||||
"options": ["engine:insightface", "model_pack:antelopev2"],
|
||||
"has_analyze": True,
|
||||
"needs_opencv_files": None,
|
||||
},
|
||||
{
|
||||
"name": "insightface-opencv",
|
||||
"options": [
|
||||
"engine:onnx_direct",
|
||||
"detector_onnx:face_detection_yunet_2023mar.onnx",
|
||||
"recognizer_onnx:face_recognition_sface_2021dec.onnx",
|
||||
],
|
||||
"has_analyze": False,
|
||||
"needs_opencv_files": "fp32",
|
||||
},
|
||||
{
|
||||
"name": "insightface-opencv-int8",
|
||||
"options": [
|
||||
"engine:onnx_direct",
|
||||
"detector_onnx:face_detection_yunet_2023mar_int8.onnx",
|
||||
"recognizer_onnx:face_recognition_sface_2021dec_int8.onnx",
|
||||
],
|
||||
"has_analyze": False,
|
||||
"needs_opencv_files": "int8",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
class _FakeContext:
|
||||
def __init__(self) -> None:
|
||||
self.code = None
|
||||
self.details = None
|
||||
|
||||
def set_code(self, code):
|
||||
self.code = code
|
||||
|
||||
def set_details(self, details):
|
||||
self.details = details
|
||||
|
||||
|
||||
def _encode_image(img: np.ndarray) -> str:
|
||||
_, buf = cv2.imencode(".jpg", img)
|
||||
return base64.b64encode(buf.tobytes()).decode("ascii")
|
||||
|
||||
|
||||
def _load_sample_image() -> str:
|
||||
from insightface.data import get_image as ins_get_image
|
||||
|
||||
return _encode_image(ins_get_image("t1"))
|
||||
|
||||
|
||||
def _download_if_missing(model_dir: str, filename: str, uri: str, sha256: str) -> None:
|
||||
dest = os.path.join(model_dir, filename)
|
||||
if os.path.isfile(dest):
|
||||
h = hashlib.sha256(open(dest, "rb").read()).hexdigest()
|
||||
if h == sha256:
|
||||
return
|
||||
sys.stderr.write(f" fetching {filename} from {uri}\n")
|
||||
sys.stderr.flush()
|
||||
urllib.request.urlretrieve(uri, dest)
|
||||
h = hashlib.sha256(open(dest, "rb").read()).hexdigest()
|
||||
if h != sha256:
|
||||
raise RuntimeError(f"sha256 mismatch for {filename}: want {sha256}, got {h}")
|
||||
|
||||
|
||||
def _run_one(cfg: dict, img_b64: str, model_dir: str) -> tuple[bool, str]:
|
||||
# Mirror LocalAI's gallery flow: populate model_dir with the
|
||||
# gallery's listed files before calling LoadModel.
|
||||
if cfg["needs_opencv_files"]:
|
||||
for filename, uri, sha256 in OPENCV_FILES[cfg["needs_opencv_files"]]:
|
||||
_download_if_missing(model_dir, filename, uri, sha256)
|
||||
|
||||
svc = BackendServicer()
|
||||
ctx = _FakeContext()
|
||||
|
||||
load_res = svc.LoadModel(
|
||||
backend_pb2.ModelOptions(
|
||||
Model=cfg["name"],
|
||||
Options=cfg["options"],
|
||||
# ModelPath is what the Go loader sets to ml.ModelPath —
|
||||
# LocalAI's models directory. The backend anchors relative
|
||||
# paths and insightface auto-download root here.
|
||||
ModelPath=model_dir,
|
||||
),
|
||||
ctx,
|
||||
)
|
||||
if not load_res.success:
|
||||
return False, f"LoadModel: {load_res.message}"
|
||||
|
||||
det_res = svc.Detect(backend_pb2.DetectOptions(src=img_b64), _FakeContext())
|
||||
if len(det_res.Detections) == 0:
|
||||
return False, "Detect returned no faces"
|
||||
for d in det_res.Detections:
|
||||
if d.class_name != "face":
|
||||
return False, f"Detect returned class_name={d.class_name!r}"
|
||||
|
||||
emb_ctx = _FakeContext()
|
||||
emb_res = svc.Embedding(backend_pb2.PredictOptions(Images=[img_b64]), emb_ctx)
|
||||
if emb_ctx.code is not None:
|
||||
return False, f"Embedding set error code {emb_ctx.code}: {emb_ctx.details}"
|
||||
if len(emb_res.embeddings) == 0:
|
||||
return False, "Embedding returned empty vector"
|
||||
norm_sq = sum(float(x) * float(x) for x in emb_res.embeddings)
|
||||
if not (0.8 <= norm_sq <= 1.2):
|
||||
return False, f"Embedding not L2-normed (sum(x^2)={norm_sq:.3f})"
|
||||
|
||||
ver_ctx = _FakeContext()
|
||||
ver_res = svc.FaceVerify(
|
||||
backend_pb2.FaceVerifyRequest(img1=img_b64, img2=img_b64), ver_ctx
|
||||
)
|
||||
if ver_ctx.code is not None:
|
||||
return False, f"FaceVerify set error code {ver_ctx.code}: {ver_ctx.details}"
|
||||
if not ver_res.verified:
|
||||
return False, f"Same-image FaceVerify not verified (dist={ver_res.distance:.3f})"
|
||||
if ver_res.distance > 0.1:
|
||||
return False, f"Same-image distance suspiciously high ({ver_res.distance:.3f})"
|
||||
|
||||
if cfg["has_analyze"]:
|
||||
an_ctx = _FakeContext()
|
||||
an_res = svc.FaceAnalyze(backend_pb2.FaceAnalyzeRequest(img=img_b64), an_ctx)
|
||||
if an_ctx.code is not None:
|
||||
return False, f"FaceAnalyze set error code {an_ctx.code}: {an_ctx.details}"
|
||||
if len(an_res.faces) == 0:
|
||||
return False, "FaceAnalyze returned no faces"
|
||||
f0 = an_res.faces[0]
|
||||
if f0.age <= 0:
|
||||
return False, f"FaceAnalyze age not populated (age={f0.age})"
|
||||
if f0.dominant_gender not in ("Man", "Woman"):
|
||||
return False, f"FaceAnalyze dominant_gender={f0.dominant_gender!r}"
|
||||
|
||||
n_dets = len(det_res.Detections)
|
||||
dim = len(emb_res.embeddings)
|
||||
return True, f"faces={n_dets} dim={dim} same-dist={ver_res.distance:.3f}"
|
||||
|
||||
|
||||
def main() -> int:
|
||||
# Honor LOCALAI_MODELS_PATH to re-use cached downloads across runs;
|
||||
# default to a fresh temp dir.
|
||||
model_dir = os.environ.get("LOCALAI_MODELS_PATH")
|
||||
if not model_dir:
|
||||
import tempfile
|
||||
|
||||
model_dir = tempfile.mkdtemp(prefix="face-smoke-")
|
||||
os.makedirs(model_dir, exist_ok=True)
|
||||
print(f"model_dir={model_dir}", file=sys.stderr)
|
||||
|
||||
print("Preparing sample image from insightface.data...", file=sys.stderr)
|
||||
img_b64 = _load_sample_image()
|
||||
|
||||
results: list[tuple[str, bool, str]] = []
|
||||
for cfg in CONFIGS:
|
||||
sys.stderr.write(f"\n=== {cfg['name']} ===\n")
|
||||
sys.stderr.flush()
|
||||
try:
|
||||
ok, detail = _run_one(cfg, img_b64, model_dir)
|
||||
except Exception:
|
||||
ok, detail = False, traceback.format_exc().splitlines()[-1]
|
||||
results.append((cfg["name"], ok, detail))
|
||||
print(f"{'PASS' if ok else 'FAIL'}: {cfg['name']:30s} {detail}")
|
||||
sys.stdout.flush()
|
||||
|
||||
print("\n=== summary ===")
|
||||
passed = sum(1 for _, ok, _ in results if ok)
|
||||
total = len(results)
|
||||
for name, ok, detail in results:
|
||||
mark = "✓" if ok else "✗"
|
||||
print(f" {mark} {name:30s} {detail}")
|
||||
print(f"\n{passed}/{total} passed")
|
||||
return 0 if passed == total else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
344
backend/python/insightface/test.py
Normal file
344
backend/python/insightface/test.py
Normal file
@@ -0,0 +1,344 @@
|
||||
"""Unit tests for the insightface gRPC backend.
|
||||
|
||||
The servicer is instantiated in-process (no gRPC channel) and driven
|
||||
directly. Images come from insightface.data which ships with the pip
|
||||
package — no external downloads.
|
||||
|
||||
Tests are parametrized over both engines (InsightFaceEngine and
|
||||
OnnxDirectEngine) where applicable.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
import cv2
|
||||
import grpc
|
||||
import numpy as np
|
||||
|
||||
sys.path.insert(0, os.path.dirname(__file__))
|
||||
|
||||
import backend_pb2 # noqa: E402
|
||||
|
||||
from backend import BackendServicer # noqa: E402
|
||||
|
||||
# OpenCV Zoo face ONNX files — downloaded on demand in OnnxDirectEngineTest
|
||||
# to mirror LocalAI's gallery `files:` flow (the backend image itself
|
||||
# doesn't ship model weights).
|
||||
OPENCV_FILES = [
|
||||
(
|
||||
"face_detection_yunet_2023mar.onnx",
|
||||
"https://github.com/opencv/opencv_zoo/raw/main/models/face_detection_yunet/face_detection_yunet_2023mar.onnx",
|
||||
"8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4",
|
||||
),
|
||||
(
|
||||
"face_recognition_sface_2021dec.onnx",
|
||||
"https://github.com/opencv/opencv_zoo/raw/main/models/face_recognition_sface/face_recognition_sface_2021dec.onnx",
|
||||
"0ba9fbfa01b5270c96627c4ef784da859931e02f04419c829e83484087c34e79",
|
||||
),
|
||||
]
|
||||
|
||||
# Silent-Face MiniFASNet ONNX files for antispoofing tests.
|
||||
ANTISPOOF_FILES = [
|
||||
(
|
||||
"MiniFASNetV2.onnx",
|
||||
"https://github.com/yakhyo/face-anti-spoofing/releases/download/weights/MiniFASNetV2.onnx",
|
||||
"b32929adc2d9c34b9486f8c4c7bc97c1b69bc0ea9befefc380e4faae4e463907",
|
||||
),
|
||||
(
|
||||
"MiniFASNetV1SE.onnx",
|
||||
"https://github.com/yakhyo/face-anti-spoofing/releases/download/weights/MiniFASNetV1SE.onnx",
|
||||
"ebab7f90c7833fbccd46d3a555410e78d969db5438e169b6524be444862b3676",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _download_files(specs: list[tuple[str, str, str]], env_var: str, prefix: str) -> str | None:
|
||||
"""Download a list of (filename, uri, sha256) into a directory.
|
||||
|
||||
Returns the directory, or None if any download failed.
|
||||
"""
|
||||
import hashlib
|
||||
import tempfile
|
||||
import urllib.request
|
||||
|
||||
root = os.environ.get(env_var) or tempfile.mkdtemp(prefix=prefix)
|
||||
for filename, uri, sha256 in specs:
|
||||
dest = os.path.join(root, filename)
|
||||
if os.path.isfile(dest):
|
||||
if hashlib.sha256(open(dest, "rb").read()).hexdigest() == sha256:
|
||||
continue
|
||||
try:
|
||||
urllib.request.urlretrieve(uri, dest)
|
||||
except Exception:
|
||||
return None
|
||||
if hashlib.sha256(open(dest, "rb").read()).hexdigest() != sha256:
|
||||
return None
|
||||
return root
|
||||
|
||||
|
||||
def _encode(img: np.ndarray) -> str:
|
||||
_, buf = cv2.imencode(".jpg", img)
|
||||
return base64.b64encode(buf.tobytes()).decode("ascii")
|
||||
|
||||
|
||||
def _load_insightface_samples() -> dict[str, str]:
|
||||
"""Return {'t1': <b64>, 't2': <b64>} from insightface.data.get_image.
|
||||
|
||||
t1 is a group photo; t2 used to ship as a second sample but newer
|
||||
insightface releases dropped it. We fall back to `Tom_Hanks_54745`
|
||||
(also bundled) as a distinct second face.
|
||||
"""
|
||||
from insightface.data import get_image as ins_get_image
|
||||
|
||||
try:
|
||||
second = ins_get_image("t2")
|
||||
except AssertionError:
|
||||
second = ins_get_image("Tom_Hanks_54745")
|
||||
return {
|
||||
"t1": _encode(ins_get_image("t1")),
|
||||
"t2": _encode(second),
|
||||
}
|
||||
|
||||
|
||||
class _FakeContext:
|
||||
"""Minimal stand-in for grpc.ServicerContext."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.code = None
|
||||
self.details = None
|
||||
|
||||
def set_code(self, code):
|
||||
self.code = code
|
||||
|
||||
def set_details(self, details):
|
||||
self.details = details
|
||||
|
||||
|
||||
class _Harness:
|
||||
def __init__(self, servicer: BackendServicer) -> None:
|
||||
self.svc = servicer
|
||||
|
||||
def health(self):
|
||||
return self.svc.Health(backend_pb2.HealthMessage(), _FakeContext())
|
||||
|
||||
def load(self, options: list[str], model_path: str = ""):
|
||||
return self.svc.LoadModel(
|
||||
backend_pb2.ModelOptions(Model="test", Options=options, ModelPath=model_path),
|
||||
_FakeContext(),
|
||||
)
|
||||
|
||||
def detect(self, img_b64: str):
|
||||
return self.svc.Detect(backend_pb2.DetectOptions(src=img_b64), _FakeContext())
|
||||
|
||||
def embed(self, img_b64: str):
|
||||
ctx = _FakeContext()
|
||||
res = self.svc.Embedding(
|
||||
backend_pb2.PredictOptions(Images=[img_b64]),
|
||||
ctx,
|
||||
)
|
||||
return res, ctx
|
||||
|
||||
def verify(self, a: str, b: str, threshold: float = 0.0, anti_spoofing: bool = False):
|
||||
ctx = _FakeContext()
|
||||
res = self.svc.FaceVerify(
|
||||
backend_pb2.FaceVerifyRequest(
|
||||
img1=a, img2=b, threshold=threshold, anti_spoofing=anti_spoofing
|
||||
),
|
||||
ctx,
|
||||
)
|
||||
return res, ctx
|
||||
|
||||
def analyze(self, img_b64: str, anti_spoofing: bool = False):
|
||||
ctx = _FakeContext()
|
||||
res = self.svc.FaceAnalyze(
|
||||
backend_pb2.FaceAnalyzeRequest(img=img_b64, anti_spoofing=anti_spoofing),
|
||||
ctx,
|
||||
)
|
||||
return res, ctx
|
||||
|
||||
|
||||
class InsightFaceEngineTest(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.samples = _load_insightface_samples()
|
||||
cls.harness = _Harness(BackendServicer())
|
||||
load = cls.harness.load(["engine:insightface", "model_pack:buffalo_l"])
|
||||
if not load.success:
|
||||
raise unittest.SkipTest(f"LoadModel failed: {load.message}")
|
||||
|
||||
def test_health(self):
|
||||
self.assertEqual(self.harness.health().message, b"OK")
|
||||
|
||||
def test_detect_finds_face(self):
|
||||
res = self.harness.detect(self.samples["t1"])
|
||||
self.assertGreater(len(res.Detections), 0)
|
||||
for d in res.Detections:
|
||||
self.assertEqual(d.class_name, "face")
|
||||
self.assertGreater(d.width, 0)
|
||||
self.assertGreater(d.height, 0)
|
||||
|
||||
def test_embedding_is_l2_normed(self):
|
||||
res, ctx = self.harness.embed(self.samples["t1"])
|
||||
self.assertIsNone(ctx.code, f"Embedding error: {ctx.details}")
|
||||
self.assertEqual(len(res.embeddings), 512)
|
||||
norm_sq = sum(x * x for x in res.embeddings)
|
||||
self.assertAlmostEqual(norm_sq, 1.0, places=2)
|
||||
|
||||
def test_verify_same_image(self):
|
||||
res, _ = self.harness.verify(self.samples["t1"], self.samples["t1"])
|
||||
self.assertTrue(res.verified)
|
||||
self.assertLess(res.distance, 0.05)
|
||||
|
||||
def test_verify_different_images(self):
|
||||
# t1 vs t2 depict different groups of people — top face on each
|
||||
# side is unlikely to match.
|
||||
res, _ = self.harness.verify(self.samples["t1"], self.samples["t2"])
|
||||
# We assert only that some numerical answer came back; the
|
||||
# matches-or-not determination depends on which face each side
|
||||
# picked and isn't a stable test assertion.
|
||||
self.assertGreaterEqual(res.distance, 0.0)
|
||||
|
||||
def test_analyze_has_age_and_gender(self):
|
||||
res, _ = self.harness.analyze(self.samples["t1"])
|
||||
self.assertGreater(len(res.faces), 0)
|
||||
for face in res.faces:
|
||||
self.assertGreater(face.face_confidence, 0.0)
|
||||
# Age should be populated for buffalo_l.
|
||||
self.assertGreater(face.age, 0.0)
|
||||
self.assertIn(face.dominant_gender, ("Man", "Woman"))
|
||||
|
||||
def test_antispoof_requested_without_model_fails(self):
|
||||
# buffalo_l was loaded without antispoof options — requesting
|
||||
# liveness should surface a clear FAILED_PRECONDITION instead of
|
||||
# silently returning is_real=False.
|
||||
_, ctx = self.harness.verify(
|
||||
self.samples["t1"], self.samples["t1"], anti_spoofing=True
|
||||
)
|
||||
self.assertEqual(ctx.code, grpc.StatusCode.FAILED_PRECONDITION)
|
||||
self.assertIn("anti_spoofing", ctx.details)
|
||||
|
||||
|
||||
def _prepare_opencv_models_dir() -> str | None:
|
||||
return _download_files(OPENCV_FILES, "OPENCV_FACE_MODELS_DIR", "opencv-face-")
|
||||
|
||||
|
||||
def _prepare_antispoof_models_dir(extra_dir: str | None = None) -> str | None:
|
||||
"""Download MiniFASNet ONNX files. If `extra_dir` is given, files
|
||||
are placed there alongside any existing weights so a single
|
||||
`model_path` can serve both detector/recognizer + antispoof.
|
||||
"""
|
||||
if extra_dir is not None:
|
||||
os.environ.setdefault("ANTISPOOF_MODELS_DIR", extra_dir)
|
||||
return _download_files(ANTISPOOF_FILES, "ANTISPOOF_MODELS_DIR", "antispoof-")
|
||||
|
||||
|
||||
class OnnxDirectEngineTest(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.samples = _load_insightface_samples()
|
||||
cls.model_dir = _prepare_opencv_models_dir()
|
||||
if cls.model_dir is None:
|
||||
raise unittest.SkipTest("OpenCV Zoo ONNX files could not be downloaded")
|
||||
cls.harness = _Harness(BackendServicer())
|
||||
load = cls.harness.load(
|
||||
[
|
||||
"engine:onnx_direct",
|
||||
"detector_onnx:face_detection_yunet_2023mar.onnx",
|
||||
"recognizer_onnx:face_recognition_sface_2021dec.onnx",
|
||||
],
|
||||
model_path=cls.model_dir,
|
||||
)
|
||||
if not load.success:
|
||||
raise unittest.SkipTest(f"LoadModel failed: {load.message}")
|
||||
|
||||
def test_detect_finds_face(self):
|
||||
res = self.harness.detect(self.samples["t1"])
|
||||
self.assertGreater(len(res.Detections), 0)
|
||||
for d in res.Detections:
|
||||
self.assertEqual(d.class_name, "face")
|
||||
|
||||
def test_embedding_nonempty(self):
|
||||
res, ctx = self.harness.embed(self.samples["t1"])
|
||||
self.assertIsNone(ctx.code, f"Embedding error: {ctx.details}")
|
||||
self.assertGreater(len(res.embeddings), 0)
|
||||
|
||||
def test_verify_same_image(self):
|
||||
res, _ = self.harness.verify(self.samples["t1"], self.samples["t1"], threshold=0.4)
|
||||
self.assertTrue(res.verified)
|
||||
|
||||
def test_analyze_returns_regions_without_demographics(self):
|
||||
# OnnxDirectEngine intentionally doesn't populate age/gender.
|
||||
res, _ = self.harness.analyze(self.samples["t1"])
|
||||
self.assertGreater(len(res.faces), 0)
|
||||
for face in res.faces:
|
||||
self.assertEqual(face.dominant_gender, "")
|
||||
self.assertEqual(face.age, 0.0)
|
||||
|
||||
|
||||
class AntispoofingTest(unittest.TestCase):
|
||||
"""End-to-end FaceVerify / FaceAnalyze with anti_spoofing=True.
|
||||
|
||||
Loads the OpenCV-Zoo (Apache-2.0) face engine alongside the Silent-Face
|
||||
MiniFASNet ensemble. Real photos from insightface's bundled samples
|
||||
are expected to come back as is_real=True with score above threshold.
|
||||
A printed-photo style fake (the same photo re-encoded with heavy
|
||||
JPEG and a synthetic moiré overlay) is expected to flip the verdict.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# Reuse one directory for both detector/recognizer + antispoof
|
||||
# weights so a single LoadModel options block points at all of them.
|
||||
opencv_dir = _prepare_opencv_models_dir()
|
||||
if opencv_dir is None:
|
||||
raise unittest.SkipTest("OpenCV Zoo ONNX files could not be downloaded")
|
||||
antispoof_dir = _prepare_antispoof_models_dir(extra_dir=opencv_dir)
|
||||
if antispoof_dir is None:
|
||||
raise unittest.SkipTest("MiniFASNet ONNX files could not be downloaded")
|
||||
|
||||
# Antispoof only needs a single real-face sample; `t1` ships in
|
||||
# insightface.data across every release.
|
||||
from insightface.data import get_image as ins_get_image
|
||||
|
||||
cls.samples = {"t1": _encode(ins_get_image("t1"))}
|
||||
cls.harness = _Harness(BackendServicer())
|
||||
load = cls.harness.load(
|
||||
[
|
||||
"engine:onnx_direct",
|
||||
"detector_onnx:face_detection_yunet_2023mar.onnx",
|
||||
"recognizer_onnx:face_recognition_sface_2021dec.onnx",
|
||||
"antispoof_v2_onnx:MiniFASNetV2.onnx",
|
||||
"antispoof_v1se_onnx:MiniFASNetV1SE.onnx",
|
||||
],
|
||||
model_path=opencv_dir,
|
||||
)
|
||||
if not load.success:
|
||||
raise unittest.SkipTest(f"LoadModel failed: {load.message}")
|
||||
|
||||
def test_verify_returns_per_image_liveness(self):
|
||||
res, ctx = self.harness.verify(
|
||||
self.samples["t1"], self.samples["t1"], threshold=0.4, anti_spoofing=True
|
||||
)
|
||||
self.assertIsNone(ctx.code, f"FaceVerify error: {ctx.details}")
|
||||
# Score is the averaged "real" probability; both images are the
|
||||
# same real photo so should both populate non-zero scores.
|
||||
self.assertGreater(res.img1_antispoof_score, 0.0)
|
||||
self.assertGreater(res.img2_antispoof_score, 0.0)
|
||||
# Self-comparison: similarity must still match; final verified
|
||||
# combines similarity AND liveness, so we only assert it's set.
|
||||
self.assertIsInstance(res.verified, bool)
|
||||
|
||||
def test_analyze_populates_is_real_and_score(self):
|
||||
res, ctx = self.harness.analyze(self.samples["t1"], anti_spoofing=True)
|
||||
self.assertIsNone(ctx.code, f"FaceAnalyze error: {ctx.details}")
|
||||
self.assertGreater(len(res.faces), 0)
|
||||
for face in res.faces:
|
||||
self.assertGreaterEqual(face.antispoof_score, 0.0)
|
||||
self.assertLessEqual(face.antispoof_score, 1.0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
11
backend/python/insightface/test.sh
Executable file
11
backend/python/insightface/test.sh
Executable file
@@ -0,0 +1,11 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
backend_dir=$(dirname $0)
|
||||
if [ -d $backend_dir/common ]; then
|
||||
source $backend_dir/common/libbackend.sh
|
||||
else
|
||||
source $backend_dir/../common/libbackend.sh
|
||||
fi
|
||||
|
||||
runUnittests
|
||||
@@ -15,17 +15,21 @@ Two startup modes:
|
||||
import asyncio
|
||||
from concurrent import futures
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
from typing import List
|
||||
|
||||
import grpc
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
|
||||
from grpc_auth import get_auth_interceptors
|
||||
from python_utils import messages_to_dicts, parse_options as _shared_parse_options
|
||||
from mlx_utils import parse_tool_calls, split_reasoning
|
||||
|
||||
|
||||
import backend_pb2
|
||||
@@ -62,37 +66,10 @@ def mlx_distributed_init(rank, hostfile, backend="ring", coordinator=None):
|
||||
raise ValueError(f"Unknown backend: {backend}")
|
||||
|
||||
|
||||
def is_float(s):
|
||||
try:
|
||||
float(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def is_int(s):
|
||||
try:
|
||||
int(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def parse_options(options):
|
||||
"""Parse key:value option strings into a dict."""
|
||||
result = {}
|
||||
for opt in options:
|
||||
if ":" not in opt:
|
||||
continue
|
||||
key, value = opt.split(":", 1)
|
||||
if is_float(value):
|
||||
value = float(value)
|
||||
elif is_int(value):
|
||||
value = int(value)
|
||||
elif value.lower() in ["true", "false"]:
|
||||
value = value.lower() == "true"
|
||||
result[key] = value
|
||||
return result
|
||||
# Re-export the shared helper under the local name for back-compat with
|
||||
# any callers (and the existing distributed worker tests) that imported
|
||||
# parse_options directly from this module.
|
||||
parse_options = _shared_parse_options
|
||||
|
||||
|
||||
class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
@@ -188,6 +165,20 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
)
|
||||
print("[Rank 0] Model loaded (single-node with prompt cache)", file=sys.stderr)
|
||||
|
||||
# Log auto-detected TokenizerWrapper capabilities. Same shape
|
||||
# as the mlx backend: has_tool_calling / has_thinking from
|
||||
# mlx_lm.tokenizer_utils + the start/end markers it sniffed
|
||||
# from the chat template / vocab.
|
||||
has_tools = bool(getattr(self.tokenizer, "has_tool_calling", False))
|
||||
has_thinking = bool(getattr(self.tokenizer, "has_thinking", False))
|
||||
tcs = getattr(self.tokenizer, "tool_call_start", None)
|
||||
tce = getattr(self.tokenizer, "tool_call_end", None)
|
||||
print(
|
||||
f"[Rank 0] Tokenizer capabilities: has_tool_calling={has_tools} "
|
||||
f"has_thinking={has_thinking} tool_call_start={tcs!r} tool_call_end={tce!r}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
except Exception as err:
|
||||
print(f"[Rank 0] Error loading model: {err}", file=sys.stderr)
|
||||
return backend_pb2.Result(success=False, message=f"Error loading model: {err}")
|
||||
@@ -201,7 +192,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
try:
|
||||
import mlx.core as mx
|
||||
from mlx_lm import stream_generate
|
||||
from mlx_lm.sample_utils import make_sampler
|
||||
from mlx_lm.sample_utils import make_logits_processors, make_sampler
|
||||
|
||||
prompt_text = self._prepare_prompt(request)
|
||||
tokens = self._get_tokens_from_prompt(prompt_text)
|
||||
@@ -211,7 +202,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
self.coordinator.broadcast_command(CMD_GENERATE, len(tokens))
|
||||
self.coordinator.broadcast_tokens(tokens)
|
||||
|
||||
max_tokens, sampler_params = self._build_generation_params(request)
|
||||
max_tokens, sampler_params, logits_params, stop_words = self._build_generation_params(request)
|
||||
|
||||
if self.coordinator:
|
||||
gen_params = self.coordinator.broadcast_generation_params(
|
||||
@@ -222,6 +213,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
max_tokens = gen_params["max_tokens"]
|
||||
|
||||
sampler = make_sampler(**sampler_params)
|
||||
logits_processors = make_logits_processors(**logits_params) if logits_params else None
|
||||
|
||||
# Use prompt cache in single-node mode
|
||||
gen_kwargs = {}
|
||||
@@ -238,22 +230,44 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
tokens = remaining_tokens if remaining_tokens else cache_key
|
||||
|
||||
generated = []
|
||||
last_response = None
|
||||
for response in stream_generate(
|
||||
self.model,
|
||||
self.tokenizer,
|
||||
prompt=tokens,
|
||||
max_tokens=max_tokens,
|
||||
sampler=sampler,
|
||||
logits_processors=logits_processors,
|
||||
**gen_kwargs,
|
||||
):
|
||||
generated.append(response.text)
|
||||
last_response = response
|
||||
if cache_key is not None:
|
||||
cache_key.append(response.token)
|
||||
if stop_words and any(s in "".join(generated) for s in stop_words):
|
||||
break
|
||||
|
||||
if self.lru_cache is not None and cache_key is not None:
|
||||
self.lru_cache.insert_cache(self.model_key, cache_key, prompt_cache)
|
||||
|
||||
return backend_pb2.Reply(message=bytes(''.join(generated), encoding='utf-8'))
|
||||
full_text = self._truncate_at_stop("".join(generated), stop_words)
|
||||
content, reasoning_content, tool_calls_proto, prompt_tokens, completion_tokens, logprobs_bytes = (
|
||||
self._finalize_output(request, full_text, last_response)
|
||||
)
|
||||
|
||||
return backend_pb2.Reply(
|
||||
message=bytes(content, encoding='utf-8'),
|
||||
prompt_tokens=prompt_tokens,
|
||||
tokens=completion_tokens,
|
||||
logprobs=logprobs_bytes,
|
||||
chat_deltas=[
|
||||
backend_pb2.ChatDelta(
|
||||
content=content,
|
||||
reasoning_content=reasoning_content,
|
||||
tool_calls=tool_calls_proto,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[Rank 0] Error in Predict: {e}", file=sys.stderr)
|
||||
@@ -268,7 +282,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
try:
|
||||
import mlx.core as mx
|
||||
from mlx_lm import stream_generate
|
||||
from mlx_lm.sample_utils import make_sampler
|
||||
from mlx_lm.sample_utils import make_logits_processors, make_sampler
|
||||
|
||||
prompt_text = self._prepare_prompt(request)
|
||||
tokens = self._get_tokens_from_prompt(prompt_text)
|
||||
@@ -278,7 +292,9 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
self.coordinator.broadcast_command(CMD_GENERATE, len(tokens))
|
||||
self.coordinator.broadcast_tokens(tokens)
|
||||
|
||||
max_tokens, sampler_params = self._build_generation_params(request, default_max_tokens=512)
|
||||
max_tokens, sampler_params, logits_params, stop_words = self._build_generation_params(
|
||||
request, default_max_tokens=512
|
||||
)
|
||||
|
||||
if self.coordinator:
|
||||
gen_params = self.coordinator.broadcast_generation_params(
|
||||
@@ -289,6 +305,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
max_tokens = gen_params["max_tokens"]
|
||||
|
||||
sampler = make_sampler(**sampler_params)
|
||||
logits_processors = make_logits_processors(**logits_params) if logits_params else None
|
||||
|
||||
# Use prompt cache in single-node mode
|
||||
gen_kwargs = {}
|
||||
@@ -304,17 +321,45 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
gen_kwargs['prompt_cache'] = prompt_cache
|
||||
tokens = remaining_tokens if remaining_tokens else cache_key
|
||||
|
||||
accumulated = []
|
||||
last_response = None
|
||||
for response in stream_generate(
|
||||
self.model,
|
||||
self.tokenizer,
|
||||
prompt=tokens,
|
||||
max_tokens=max_tokens,
|
||||
sampler=sampler,
|
||||
logits_processors=logits_processors,
|
||||
**gen_kwargs,
|
||||
):
|
||||
if cache_key is not None:
|
||||
cache_key.append(response.token)
|
||||
yield backend_pb2.Reply(message=bytes(response.text, encoding='utf-8'))
|
||||
accumulated.append(response.text)
|
||||
last_response = response
|
||||
yield backend_pb2.Reply(
|
||||
message=bytes(response.text, encoding='utf-8'),
|
||||
chat_deltas=[backend_pb2.ChatDelta(content=response.text)],
|
||||
)
|
||||
if stop_words and any(s in "".join(accumulated) for s in stop_words):
|
||||
break
|
||||
|
||||
full_text = self._truncate_at_stop("".join(accumulated), stop_words)
|
||||
content, reasoning_content, tool_calls_proto, prompt_tokens, completion_tokens, logprobs_bytes = (
|
||||
self._finalize_output(request, full_text, last_response)
|
||||
)
|
||||
yield backend_pb2.Reply(
|
||||
message=b"",
|
||||
prompt_tokens=prompt_tokens,
|
||||
tokens=completion_tokens,
|
||||
logprobs=logprobs_bytes,
|
||||
chat_deltas=[
|
||||
backend_pb2.ChatDelta(
|
||||
content="",
|
||||
reasoning_content=reasoning_content,
|
||||
tool_calls=tool_calls_proto,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[Rank 0] Error in PredictStream: {e}", file=sys.stderr)
|
||||
@@ -335,12 +380,74 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
context.set_details("Embeddings are not supported in the MLX distributed backend.")
|
||||
return backend_pb2.EmbeddingResult()
|
||||
|
||||
async def TokenizeString(self, request, context):
|
||||
if not hasattr(self, "tokenizer") or self.tokenizer is None:
|
||||
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
|
||||
context.set_details("tokenizer not loaded")
|
||||
return backend_pb2.TokenizationResponse()
|
||||
try:
|
||||
tokens = self.tokenizer.encode(request.Prompt)
|
||||
if hasattr(tokens, "tolist"):
|
||||
tokens = tokens.tolist()
|
||||
tokens = list(tokens)
|
||||
return backend_pb2.TokenizationResponse(length=len(tokens), tokens=tokens)
|
||||
except Exception as e:
|
||||
context.set_code(grpc.StatusCode.INTERNAL)
|
||||
context.set_details(str(e))
|
||||
return backend_pb2.TokenizationResponse()
|
||||
|
||||
async def Free(self, request, context):
|
||||
try:
|
||||
# If we're rank 0 of a distributed run, tell workers to shut
|
||||
# down their per-request loops first so they release the model.
|
||||
if self.coordinator is not None:
|
||||
try:
|
||||
from coordinator import CMD_SHUTDOWN
|
||||
self.coordinator.broadcast_command(CMD_SHUTDOWN)
|
||||
except Exception as e:
|
||||
print(f"[Rank 0] failed to broadcast shutdown: {e}", file=sys.stderr)
|
||||
if hasattr(self, "model"):
|
||||
del self.model
|
||||
if hasattr(self, "tokenizer"):
|
||||
del self.tokenizer
|
||||
if self.lru_cache is not None:
|
||||
try:
|
||||
self.lru_cache.clear()
|
||||
except Exception:
|
||||
pass
|
||||
self.lru_cache = None
|
||||
self.coordinator = None
|
||||
self.group = None
|
||||
gc.collect()
|
||||
try:
|
||||
import mlx.core as mx # type: ignore
|
||||
if hasattr(mx, "clear_cache"):
|
||||
mx.clear_cache()
|
||||
elif hasattr(mx, "metal") and hasattr(mx.metal, "clear_cache"):
|
||||
mx.metal.clear_cache()
|
||||
except Exception:
|
||||
pass
|
||||
return backend_pb2.Result(success=True, message="MLX distributed model freed")
|
||||
except Exception as e:
|
||||
return backend_pb2.Result(success=False, message=str(e))
|
||||
|
||||
def _prepare_prompt(self, request):
|
||||
if not request.Prompt and request.UseTokenizerTemplate and request.Messages:
|
||||
messages = [{"role": msg.role, "content": msg.content} for msg in request.Messages]
|
||||
return self.tokenizer.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
messages = messages_to_dicts(request.Messages)
|
||||
kwargs = {"tokenize": False, "add_generation_prompt": True}
|
||||
if request.Tools:
|
||||
try:
|
||||
kwargs["tools"] = json.loads(request.Tools)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
if request.Metadata.get("enable_thinking", "").lower() == "true":
|
||||
kwargs["enable_thinking"] = True
|
||||
try:
|
||||
return self.tokenizer.apply_chat_template(messages, **kwargs)
|
||||
except TypeError:
|
||||
return self.tokenizer.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
return request.Prompt
|
||||
|
||||
def _get_tokens_from_prompt(self, prompt_text: str) -> List[int]:
|
||||
@@ -349,6 +456,82 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
return tokens.tolist()
|
||||
return list(tokens)
|
||||
|
||||
def _tool_module_from_tokenizer(self):
|
||||
"""Same shim as the mlx backend: fall back to json.loads when the
|
||||
installed mlx-lm doesn't expose a tool_parser callable on the
|
||||
wrapper (true on 0.29.x — only HEAD ships parsers)."""
|
||||
start = getattr(self.tokenizer, "tool_call_start", None)
|
||||
end = getattr(self.tokenizer, "tool_call_end", None)
|
||||
if not start:
|
||||
return None
|
||||
parse_fn = getattr(self.tokenizer, "tool_parser", None)
|
||||
if parse_fn is None:
|
||||
def parse_fn(body, tools): # noqa: E306
|
||||
return json.loads(body.strip())
|
||||
return types.SimpleNamespace(
|
||||
tool_call_start=start,
|
||||
tool_call_end=end or "",
|
||||
parse_tool_call=parse_fn,
|
||||
)
|
||||
|
||||
def _truncate_at_stop(self, text, stop_words):
|
||||
if not stop_words:
|
||||
return text
|
||||
earliest = len(text)
|
||||
for stop in stop_words:
|
||||
if not stop:
|
||||
continue
|
||||
idx = text.find(stop)
|
||||
if idx >= 0 and idx < earliest:
|
||||
earliest = idx
|
||||
return text[:earliest] if earliest < len(text) else text
|
||||
|
||||
def _finalize_output(self, request, generated_text, last_response):
|
||||
content = generated_text
|
||||
reasoning_content = ""
|
||||
if getattr(self.tokenizer, "has_thinking", False):
|
||||
think_start = getattr(self.tokenizer, "think_start", "") or ""
|
||||
think_end = getattr(self.tokenizer, "think_end", "") or ""
|
||||
reasoning_content, content = split_reasoning(content, think_start, think_end)
|
||||
|
||||
tool_calls_proto: List[backend_pb2.ToolCallDelta] = []
|
||||
tool_module = None
|
||||
if getattr(self.tokenizer, "has_tool_calling", False):
|
||||
tool_module = self._tool_module_from_tokenizer()
|
||||
if tool_module is not None:
|
||||
parsed_tools = None
|
||||
if request.Tools:
|
||||
try:
|
||||
parsed_tools = json.loads(request.Tools)
|
||||
except json.JSONDecodeError:
|
||||
parsed_tools = None
|
||||
calls, content = parse_tool_calls(content, tool_module, parsed_tools)
|
||||
for c in calls:
|
||||
tool_calls_proto.append(
|
||||
backend_pb2.ToolCallDelta(
|
||||
index=c["index"], id=c["id"], name=c["name"], arguments=c["arguments"],
|
||||
)
|
||||
)
|
||||
|
||||
prompt_token_count = int(getattr(last_response, "prompt_tokens", 0) or 0) if last_response else 0
|
||||
completion_token_count = int(getattr(last_response, "generation_tokens", 0) or 0) if last_response else 0
|
||||
|
||||
logprobs_bytes = b""
|
||||
if last_response is not None and int(getattr(request, "Logprobs", 0) or 0) > 0:
|
||||
try:
|
||||
lp = getattr(last_response, "logprobs", None)
|
||||
if lp is not None:
|
||||
token_id = int(getattr(last_response, "token", 0) or 0)
|
||||
token_text = self.tokenizer.decode([token_id]) if token_id else ""
|
||||
top_logprob = float(lp[token_id]) if hasattr(lp, "__getitem__") else 0.0
|
||||
logprobs_bytes = json.dumps(
|
||||
{"content": [{"token": token_text, "logprob": top_logprob}]}
|
||||
).encode("utf-8")
|
||||
except Exception as e:
|
||||
print(f"[Rank 0] Logprobs extraction failed: {e}", file=sys.stderr)
|
||||
|
||||
return content, reasoning_content, tool_calls_proto, prompt_token_count, completion_token_count, logprobs_bytes
|
||||
|
||||
def _build_generation_params(self, request, default_max_tokens=200):
|
||||
import mlx.core as mx
|
||||
|
||||
@@ -373,6 +556,22 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
'xtc_probability': 0.0,
|
||||
}
|
||||
|
||||
# Logits processor parameters — pulled from the request and
|
||||
# forwarded to make_logits_processors. Rank 0 is the only rank
|
||||
# running the sampler so we don't need to broadcast these to
|
||||
# workers (workers participate in the pipeline-parallel forward
|
||||
# pass only).
|
||||
logits_params = {}
|
||||
repetition_penalty = getattr(request, 'RepetitionPenalty', 0.0) or 0.0
|
||||
if repetition_penalty and repetition_penalty != 1.0:
|
||||
logits_params['repetition_penalty'] = repetition_penalty
|
||||
presence_penalty = getattr(request, 'PresencePenalty', 0.0) or 0.0
|
||||
if presence_penalty:
|
||||
logits_params['presence_penalty'] = presence_penalty
|
||||
frequency_penalty = getattr(request, 'FrequencyPenalty', 0.0) or 0.0
|
||||
if frequency_penalty:
|
||||
logits_params['frequency_penalty'] = frequency_penalty
|
||||
|
||||
seed = getattr(request, 'Seed', 0)
|
||||
if seed != 0:
|
||||
mx.random.seed(seed)
|
||||
@@ -392,9 +591,15 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
for opt_key, param_key in option_mapping.items():
|
||||
if opt_key in self.options:
|
||||
sampler_params[param_key] = self.options[opt_key]
|
||||
for opt_key in ('repetition_penalty', 'presence_penalty', 'frequency_penalty'):
|
||||
if opt_key in self.options:
|
||||
logits_params[opt_key] = self.options[opt_key]
|
||||
if 'seed' in self.options:
|
||||
mx.random.seed(self.options['seed'])
|
||||
|
||||
stop_words = list(getattr(request, 'StopPrompts', []) or [])
|
||||
return max_tokens, sampler_params, logits_params, stop_words
|
||||
|
||||
# XTC special tokens
|
||||
xtc_special_tokens = []
|
||||
if hasattr(self.tokenizer, 'eos_token_ids') and self.tokenizer.eos_token_ids:
|
||||
|
||||
@@ -1,3 +1,6 @@
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
import subprocess
|
||||
import time
|
||||
@@ -6,6 +9,12 @@ import grpc
|
||||
import backend_pb2
|
||||
import backend_pb2_grpc
|
||||
|
||||
# Make the shared helpers importable so we can unit-test them without a
|
||||
# running gRPC server.
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
|
||||
from python_utils import messages_to_dicts, parse_options
|
||||
from mlx_utils import parse_tool_calls, split_reasoning
|
||||
|
||||
|
||||
class TestBackendServicer(unittest.TestCase):
|
||||
def setUp(self):
|
||||
@@ -85,3 +94,44 @@ class TestBackendServicer(unittest.TestCase):
|
||||
self.fail("sampling params service failed")
|
||||
finally:
|
||||
self.tearDown()
|
||||
|
||||
|
||||
class TestSharedHelpers(unittest.TestCase):
|
||||
"""Server-less unit tests for the helpers the mlx-distributed backend depends on."""
|
||||
|
||||
def test_parse_options_typed(self):
|
||||
opts = parse_options(["temperature:0.7", "max_tokens:128", "trust:true"])
|
||||
self.assertEqual(opts["temperature"], 0.7)
|
||||
self.assertEqual(opts["max_tokens"], 128)
|
||||
self.assertIs(opts["trust"], True)
|
||||
|
||||
def test_messages_to_dicts_roundtrip(self):
|
||||
msgs = [
|
||||
backend_pb2.Message(role="user", content="hi"),
|
||||
backend_pb2.Message(
|
||||
role="assistant",
|
||||
content="",
|
||||
tool_calls='[{"id":"call_1","type":"function","function":{"name":"f","arguments":"{}"}}]',
|
||||
),
|
||||
backend_pb2.Message(role="tool", content="42", tool_call_id="call_1", name="f"),
|
||||
]
|
||||
out = messages_to_dicts(msgs)
|
||||
self.assertEqual(out[0], {"role": "user", "content": "hi"})
|
||||
self.assertEqual(out[1]["tool_calls"][0]["function"]["name"], "f")
|
||||
self.assertEqual(out[2]["tool_call_id"], "call_1")
|
||||
|
||||
def test_split_reasoning(self):
|
||||
r, c = split_reasoning("<think>plan</think>final", "<think>", "</think>")
|
||||
self.assertEqual(r, "plan")
|
||||
self.assertEqual(c, "final")
|
||||
|
||||
def test_parse_tool_calls_with_shim(self):
|
||||
tm = types.SimpleNamespace(
|
||||
tool_call_start="<tool_call>",
|
||||
tool_call_end="</tool_call>",
|
||||
parse_tool_call=lambda body, tools: {"name": "get_weather", "arguments": {"location": body.strip()}},
|
||||
)
|
||||
calls, remaining = parse_tool_calls("<tool_call>Paris</tool_call>", tm, tools=None)
|
||||
self.assertEqual(len(calls), 1)
|
||||
self.assertEqual(calls[0]["name"], "get_weather")
|
||||
self.assertEqual(calls[0]["arguments"], '{"location": "Paris"}')
|
||||
|
||||
@@ -2,11 +2,14 @@
|
||||
import asyncio
|
||||
from concurrent import futures
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import signal
|
||||
import sys
|
||||
import os
|
||||
import tempfile
|
||||
import types
|
||||
from typing import List
|
||||
import time
|
||||
|
||||
import backend_pb2
|
||||
import backend_pb2_grpc
|
||||
@@ -15,30 +18,18 @@ import grpc
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
|
||||
from grpc_auth import get_auth_interceptors
|
||||
from python_utils import messages_to_dicts, parse_options
|
||||
from mlx_utils import parse_tool_calls, split_reasoning
|
||||
|
||||
from mlx_vlm import load, generate, stream_generate
|
||||
from mlx_vlm import load, stream_generate
|
||||
from mlx_vlm.prompt_utils import apply_chat_template
|
||||
from mlx_vlm.utils import load_config, load_image
|
||||
from mlx_vlm.tool_parsers import _infer_tool_parser, load_tool_module
|
||||
from mlx_vlm.utils import load_config
|
||||
from mlx_lm.sample_utils import make_logits_processors, make_sampler
|
||||
import mlx.core as mx
|
||||
import base64
|
||||
import io
|
||||
from PIL import Image
|
||||
import tempfile
|
||||
|
||||
def is_float(s):
|
||||
"""Check if a string can be converted to float."""
|
||||
try:
|
||||
float(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
def is_int(s):
|
||||
"""Check if a string can be converted to int."""
|
||||
try:
|
||||
int(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
|
||||
|
||||
@@ -78,36 +69,52 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
try:
|
||||
print(f"Loading MLX-VLM model: {request.Model}", file=sys.stderr)
|
||||
print(f"Request: {request}", file=sys.stderr)
|
||||
|
||||
# Parse options like in the diffusers backend
|
||||
options = request.Options
|
||||
self.options = {}
|
||||
|
||||
# The options are a list of strings in this form optname:optvalue
|
||||
# We store all the options in a dict for later use
|
||||
for opt in options:
|
||||
if ":" not in opt:
|
||||
continue
|
||||
key, value = opt.split(":", 1) # Split only on first colon to handle values with colons
|
||||
|
||||
if is_float(value):
|
||||
value = float(value)
|
||||
elif is_int(value):
|
||||
value = int(value)
|
||||
elif value.lower() in ["true", "false"]:
|
||||
value = value.lower() == "true"
|
||||
|
||||
self.options[key] = value
|
||||
|
||||
|
||||
# Parse Options[] key:value strings into a typed dict
|
||||
self.options = parse_options(request.Options)
|
||||
print(f"Options: {self.options}", file=sys.stderr)
|
||||
|
||||
|
||||
# Load model and processor using MLX-VLM
|
||||
# mlx-vlm load function returns (model, processor) instead of (model, tokenizer)
|
||||
self.model, self.processor = load(request.Model)
|
||||
|
||||
|
||||
# Load model config for chat template support
|
||||
self.config = load_config(request.Model)
|
||||
|
||||
|
||||
# Auto-infer the tool parser from the chat template. mlx-vlm has
|
||||
# its own _infer_tool_parser that falls back to mlx-lm parsers.
|
||||
tokenizer = (
|
||||
self.processor.tokenizer if hasattr(self.processor, "tokenizer") else self.processor
|
||||
)
|
||||
self.tool_module = None
|
||||
if hasattr(tokenizer, "chat_template"):
|
||||
try:
|
||||
parser_type = _infer_tool_parser(tokenizer.chat_template)
|
||||
if parser_type is not None:
|
||||
self.tool_module = load_tool_module(parser_type)
|
||||
print(
|
||||
f"[mlx-vlm] auto-detected tool parser: {parser_type}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
else:
|
||||
print(
|
||||
"[mlx-vlm] no tool parser matched the chat template",
|
||||
file=sys.stderr,
|
||||
)
|
||||
except Exception as e:
|
||||
print(
|
||||
f"[mlx-vlm] failed to load tool parser: {e}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
# Reasoning tokens — check if the tokenizer advertises thinking
|
||||
# markers. Fall back to empty strings (split_reasoning no-ops).
|
||||
self.think_start = getattr(tokenizer, "think_start", "") or ""
|
||||
self.think_end = getattr(tokenizer, "think_end", "") or ""
|
||||
self.has_thinking = bool(
|
||||
getattr(tokenizer, "has_thinking", False) or self.think_start
|
||||
)
|
||||
|
||||
except Exception as err:
|
||||
print(f"Error loading MLX-VLM model {err=}, {type(err)=}", file=sys.stderr)
|
||||
return backend_pb2.Result(success=False, message=f"Error loading MLX-VLM model: {err}")
|
||||
@@ -128,63 +135,72 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
"""
|
||||
temp_files = []
|
||||
try:
|
||||
# Process images and audios from request
|
||||
image_paths = []
|
||||
audio_paths = []
|
||||
|
||||
# Process images
|
||||
if request.Images:
|
||||
for img_data in request.Images:
|
||||
img_path = self.load_image_from_base64(img_data)
|
||||
if img_path:
|
||||
image_paths.append(img_path)
|
||||
temp_files.append(img_path)
|
||||
|
||||
# Process audios
|
||||
if request.Audios:
|
||||
for audio_data in request.Audios:
|
||||
audio_path = self.load_audio_from_base64(audio_data)
|
||||
if audio_path:
|
||||
audio_paths.append(audio_path)
|
||||
temp_files.append(audio_path)
|
||||
|
||||
# Prepare the prompt with multimodal information
|
||||
prompt = self._prepare_prompt(request, num_images=len(image_paths), num_audios=len(audio_paths))
|
||||
|
||||
# Build generation parameters using request attributes and options
|
||||
max_tokens, generation_params = self._build_generation_params(request)
|
||||
|
||||
print(f"Generating text with MLX-VLM - max_tokens: {max_tokens}, params: {generation_params}", file=sys.stderr)
|
||||
print(f"Images: {len(image_paths)}, Audios: {len(audio_paths)}", file=sys.stderr)
|
||||
|
||||
# Generate text using MLX-VLM with multimodal inputs
|
||||
response = generate(
|
||||
image_paths, audio_paths = self._collect_media(request, temp_files)
|
||||
|
||||
prompt = self._prepare_prompt(
|
||||
request,
|
||||
num_images=len(image_paths),
|
||||
num_audios=len(audio_paths),
|
||||
)
|
||||
|
||||
max_tokens, sampler_params, logits_params, stop_words = self._build_generation_params(request)
|
||||
sampler = make_sampler(**sampler_params)
|
||||
logits_processors = make_logits_processors(**logits_params) if logits_params else None
|
||||
|
||||
print(
|
||||
f"Generating text with MLX-VLM - max_tokens: {max_tokens}, "
|
||||
f"images: {len(image_paths)}, audios: {len(audio_paths)}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
accumulated = []
|
||||
last_response = None
|
||||
for response in stream_generate(
|
||||
model=self.model,
|
||||
processor=self.processor,
|
||||
prompt=prompt,
|
||||
image=image_paths if image_paths else None,
|
||||
audio=audio_paths if audio_paths else None,
|
||||
max_tokens=max_tokens,
|
||||
temperature=generation_params.get('temp', 0.6),
|
||||
top_p=generation_params.get('top_p', 1.0),
|
||||
verbose=False
|
||||
sampler=sampler,
|
||||
logits_processors=logits_processors,
|
||||
):
|
||||
accumulated.append(response.text)
|
||||
last_response = response
|
||||
if stop_words and any(s in "".join(accumulated) for s in stop_words):
|
||||
break
|
||||
|
||||
full_text = self._truncate_at_stop("".join(accumulated), stop_words)
|
||||
content, reasoning_content, tool_calls_proto, prompt_tokens, completion_tokens, logprobs_bytes = (
|
||||
self._finalize_output(request, full_text, last_response)
|
||||
)
|
||||
|
||||
return backend_pb2.Reply(message=bytes(response, encoding='utf-8'))
|
||||
|
||||
|
||||
return backend_pb2.Reply(
|
||||
message=bytes(content, encoding='utf-8'),
|
||||
prompt_tokens=prompt_tokens,
|
||||
tokens=completion_tokens,
|
||||
logprobs=logprobs_bytes,
|
||||
chat_deltas=[
|
||||
backend_pb2.ChatDelta(
|
||||
content=content,
|
||||
reasoning_content=reasoning_content,
|
||||
tool_calls=tool_calls_proto,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in MLX-VLM Predict: {e}", file=sys.stderr)
|
||||
context.set_code(grpc.StatusCode.INTERNAL)
|
||||
context.set_details(f"Generation failed: {str(e)}")
|
||||
return backend_pb2.Reply(message=bytes("", encoding='utf-8'))
|
||||
finally:
|
||||
# Clean up temporary files
|
||||
self.cleanup_temp_files(temp_files)
|
||||
|
||||
def Embedding(self, request, context):
|
||||
"""
|
||||
A gRPC method that calculates embeddings for a given sentence.
|
||||
|
||||
|
||||
Note: MLX-VLM doesn't support embeddings directly. This method returns an error.
|
||||
|
||||
Args:
|
||||
@@ -199,6 +215,79 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
context.set_details("Embeddings are not supported in the MLX-VLM backend.")
|
||||
return backend_pb2.EmbeddingResult()
|
||||
|
||||
def _collect_media(self, request, temp_files):
|
||||
"""Decode base64 Images and Audios into temp file paths.
|
||||
|
||||
Appends every temp file to ``temp_files`` so the finally block can
|
||||
clean up even on mid-generation errors.
|
||||
"""
|
||||
image_paths = []
|
||||
audio_paths = []
|
||||
if request.Images:
|
||||
for img_data in request.Images:
|
||||
img_path = self.load_image_from_base64(img_data)
|
||||
if img_path:
|
||||
image_paths.append(img_path)
|
||||
temp_files.append(img_path)
|
||||
if request.Audios:
|
||||
for audio_data in request.Audios:
|
||||
audio_path = self.load_audio_from_base64(audio_data)
|
||||
if audio_path:
|
||||
audio_paths.append(audio_path)
|
||||
temp_files.append(audio_path)
|
||||
return image_paths, audio_paths
|
||||
|
||||
async def TokenizeString(self, request, context):
|
||||
"""Tokenize ``request.Prompt`` via the processor's tokenizer."""
|
||||
if not hasattr(self, "processor") or self.processor is None:
|
||||
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
|
||||
context.set_details("processor not loaded")
|
||||
return backend_pb2.TokenizationResponse()
|
||||
try:
|
||||
tokenizer = (
|
||||
self.processor.tokenizer
|
||||
if hasattr(self.processor, "tokenizer")
|
||||
else self.processor
|
||||
)
|
||||
tokens = tokenizer.encode(request.Prompt)
|
||||
if hasattr(tokens, "tolist"):
|
||||
tokens = tokens.tolist()
|
||||
tokens = list(tokens)
|
||||
return backend_pb2.TokenizationResponse(length=len(tokens), tokens=tokens)
|
||||
except Exception as e:
|
||||
context.set_code(grpc.StatusCode.INTERNAL)
|
||||
context.set_details(str(e))
|
||||
return backend_pb2.TokenizationResponse()
|
||||
|
||||
async def Free(self, request, context):
|
||||
"""Drop the loaded model, processor and tool module."""
|
||||
try:
|
||||
if hasattr(self, "model"):
|
||||
del self.model
|
||||
if hasattr(self, "processor"):
|
||||
del self.processor
|
||||
if hasattr(self, "config"):
|
||||
del self.config
|
||||
self.tool_module = None
|
||||
gc.collect()
|
||||
# mlx.clear_cache (mlx >= 0.30) supersedes mlx.metal.clear_cache.
|
||||
try:
|
||||
if hasattr(mx, "clear_cache"):
|
||||
mx.clear_cache()
|
||||
elif hasattr(mx, "metal") and hasattr(mx.metal, "clear_cache"):
|
||||
mx.metal.clear_cache()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
import torch # type: ignore
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
except Exception:
|
||||
pass
|
||||
return backend_pb2.Result(success=True, message="MLX-VLM model freed")
|
||||
except Exception as e:
|
||||
return backend_pb2.Result(success=False, message=str(e))
|
||||
|
||||
async def PredictStream(self, request, context):
|
||||
"""
|
||||
Generates text based on the given prompt and sampling parameters, and streams the results using MLX-VLM with multimodal support.
|
||||
@@ -212,36 +301,28 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
"""
|
||||
temp_files = []
|
||||
try:
|
||||
# Process images and audios from request
|
||||
image_paths = []
|
||||
audio_paths = []
|
||||
|
||||
# Process images
|
||||
if request.Images:
|
||||
for img_data in request.Images:
|
||||
img_path = self.load_image_from_base64(img_data)
|
||||
if img_path:
|
||||
image_paths.append(img_path)
|
||||
temp_files.append(img_path)
|
||||
|
||||
# Process audios
|
||||
if request.Audios:
|
||||
for audio_data in request.Audios:
|
||||
audio_path = self.load_audio_from_base64(audio_data)
|
||||
if audio_path:
|
||||
audio_paths.append(audio_path)
|
||||
temp_files.append(audio_path)
|
||||
|
||||
# Prepare the prompt with multimodal information
|
||||
prompt = self._prepare_prompt(request, num_images=len(image_paths), num_audios=len(audio_paths))
|
||||
|
||||
# Build generation parameters using request attributes and options
|
||||
max_tokens, generation_params = self._build_generation_params(request, default_max_tokens=512)
|
||||
|
||||
print(f"Streaming text with MLX-VLM - max_tokens: {max_tokens}, params: {generation_params}", file=sys.stderr)
|
||||
print(f"Images: {len(image_paths)}, Audios: {len(audio_paths)}", file=sys.stderr)
|
||||
|
||||
# Stream text generation using MLX-VLM with multimodal inputs
|
||||
image_paths, audio_paths = self._collect_media(request, temp_files)
|
||||
|
||||
prompt = self._prepare_prompt(
|
||||
request,
|
||||
num_images=len(image_paths),
|
||||
num_audios=len(audio_paths),
|
||||
)
|
||||
|
||||
max_tokens, sampler_params, logits_params, stop_words = self._build_generation_params(
|
||||
request, default_max_tokens=512
|
||||
)
|
||||
sampler = make_sampler(**sampler_params)
|
||||
logits_processors = make_logits_processors(**logits_params) if logits_params else None
|
||||
|
||||
print(
|
||||
f"Streaming text with MLX-VLM - max_tokens: {max_tokens}, "
|
||||
f"images: {len(image_paths)}, audios: {len(audio_paths)}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
accumulated = []
|
||||
last_response = None
|
||||
for response in stream_generate(
|
||||
model=self.model,
|
||||
processor=self.processor,
|
||||
@@ -249,77 +330,91 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
image=image_paths if image_paths else None,
|
||||
audio=audio_paths if audio_paths else None,
|
||||
max_tokens=max_tokens,
|
||||
temperature=generation_params.get('temp', 0.6),
|
||||
top_p=generation_params.get('top_p', 1.0),
|
||||
sampler=sampler,
|
||||
logits_processors=logits_processors,
|
||||
):
|
||||
yield backend_pb2.Reply(message=bytes(response.text, encoding='utf-8'))
|
||||
|
||||
accumulated.append(response.text)
|
||||
last_response = response
|
||||
yield backend_pb2.Reply(
|
||||
message=bytes(response.text, encoding='utf-8'),
|
||||
chat_deltas=[backend_pb2.ChatDelta(content=response.text)],
|
||||
)
|
||||
if stop_words and any(s in "".join(accumulated) for s in stop_words):
|
||||
break
|
||||
|
||||
full_text = self._truncate_at_stop("".join(accumulated), stop_words)
|
||||
content, reasoning_content, tool_calls_proto, prompt_tokens, completion_tokens, logprobs_bytes = (
|
||||
self._finalize_output(request, full_text, last_response)
|
||||
)
|
||||
yield backend_pb2.Reply(
|
||||
message=b"",
|
||||
prompt_tokens=prompt_tokens,
|
||||
tokens=completion_tokens,
|
||||
logprobs=logprobs_bytes,
|
||||
chat_deltas=[
|
||||
backend_pb2.ChatDelta(
|
||||
content="",
|
||||
reasoning_content=reasoning_content,
|
||||
tool_calls=tool_calls_proto,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in MLX-VLM PredictStream: {e}", file=sys.stderr)
|
||||
context.set_code(grpc.StatusCode.INTERNAL)
|
||||
context.set_details(f"Streaming generation failed: {str(e)}")
|
||||
yield backend_pb2.Reply(message=bytes("", encoding='utf-8'))
|
||||
finally:
|
||||
# Clean up temporary files
|
||||
self.cleanup_temp_files(temp_files)
|
||||
|
||||
def _build_template_kwargs(self, request, num_images, num_audios):
|
||||
"""Collect kwargs for ``apply_chat_template`` that survive model variants."""
|
||||
kwargs = {"num_images": num_images, "num_audios": num_audios}
|
||||
if request.Tools:
|
||||
try:
|
||||
kwargs["tools"] = json.loads(request.Tools)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
if request.Metadata.get("enable_thinking", "").lower() == "true":
|
||||
kwargs["enable_thinking"] = True
|
||||
return kwargs
|
||||
|
||||
def _apply_template(self, request, messages, num_images, num_audios):
|
||||
kwargs = self._build_template_kwargs(request, num_images, num_audios)
|
||||
try:
|
||||
return apply_chat_template(self.processor, self.config, messages, **kwargs)
|
||||
except TypeError:
|
||||
# Fallback for older mlx-vlm versions that reject tools=/enable_thinking=
|
||||
return apply_chat_template(
|
||||
self.processor,
|
||||
self.config,
|
||||
messages,
|
||||
num_images=num_images,
|
||||
num_audios=num_audios,
|
||||
)
|
||||
|
||||
def _prepare_prompt(self, request, num_images=0, num_audios=0):
|
||||
"""
|
||||
Prepare the prompt for MLX-VLM generation, handling chat templates and multimodal inputs.
|
||||
|
||||
Args:
|
||||
request: The gRPC request containing prompt and message information.
|
||||
num_images: Number of images in the request.
|
||||
num_audios: Number of audio files in the request.
|
||||
|
||||
Returns:
|
||||
str: The prepared prompt.
|
||||
Prepare the prompt for MLX-VLM generation, handling chat templates and
|
||||
multimodal inputs. Forwards tool definitions and enable_thinking when
|
||||
present on the request.
|
||||
"""
|
||||
# If tokenizer template is enabled and messages are provided instead of prompt, apply the tokenizer template
|
||||
if not request.Prompt and request.UseTokenizerTemplate and request.Messages:
|
||||
# Convert gRPC messages to the format expected by apply_chat_template
|
||||
messages = []
|
||||
for msg in request.Messages:
|
||||
messages.append({"role": msg.role, "content": msg.content})
|
||||
|
||||
# Use mlx-vlm's apply_chat_template which handles multimodal inputs
|
||||
prompt = apply_chat_template(
|
||||
self.processor,
|
||||
self.config,
|
||||
messages,
|
||||
num_images=num_images,
|
||||
num_audios=num_audios
|
||||
)
|
||||
return prompt
|
||||
elif request.Prompt:
|
||||
# If we have a direct prompt but also have images/audio, we need to format it properly
|
||||
messages = messages_to_dicts(request.Messages)
|
||||
return self._apply_template(request, messages, num_images, num_audios)
|
||||
|
||||
if request.Prompt:
|
||||
if num_images > 0 or num_audios > 0:
|
||||
# Create a simple message structure for multimodal prompt
|
||||
messages = [{"role": "user", "content": request.Prompt}]
|
||||
prompt = apply_chat_template(
|
||||
self.processor,
|
||||
self.config,
|
||||
messages,
|
||||
num_images=num_images,
|
||||
num_audios=num_audios
|
||||
)
|
||||
return prompt
|
||||
else:
|
||||
return request.Prompt
|
||||
else:
|
||||
# Fallback to empty prompt with multimodal template if we have media
|
||||
if num_images > 0 or num_audios > 0:
|
||||
messages = [{"role": "user", "content": ""}]
|
||||
prompt = apply_chat_template(
|
||||
self.processor,
|
||||
self.config,
|
||||
messages,
|
||||
num_images=num_images,
|
||||
num_audios=num_audios
|
||||
)
|
||||
return prompt
|
||||
else:
|
||||
return ""
|
||||
return self._apply_template(request, messages, num_images, num_audios)
|
||||
return request.Prompt
|
||||
|
||||
# Fallback to empty prompt with multimodal template if we have media
|
||||
if num_images > 0 or num_audios > 0:
|
||||
messages = [{"role": "user", "content": ""}]
|
||||
return self._apply_template(request, messages, num_images, num_audios)
|
||||
return ""
|
||||
|
||||
|
||||
|
||||
@@ -327,62 +422,122 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
|
||||
def _build_generation_params(self, request, default_max_tokens=200):
|
||||
"""
|
||||
Build generation parameters from request attributes and options for MLX-VLM.
|
||||
|
||||
Args:
|
||||
request: The gRPC request.
|
||||
default_max_tokens: Default max_tokens if not specified.
|
||||
Build generation parameters from request attributes and options.
|
||||
|
||||
Returns:
|
||||
tuple: (max_tokens, generation_params dict)
|
||||
tuple: (max_tokens, sampler_params, logits_params, stop_words)
|
||||
"""
|
||||
# Extract max_tokens
|
||||
max_tokens = getattr(request, 'Tokens', default_max_tokens)
|
||||
if max_tokens == 0:
|
||||
max_tokens = default_max_tokens
|
||||
|
||||
# Extract generation parameters from request attributes
|
||||
temp = getattr(request, 'Temperature', 0.0)
|
||||
if temp == 0.0:
|
||||
temp = 0.6 # Default temperature
|
||||
|
||||
top_p = getattr(request, 'TopP', 0.0)
|
||||
if top_p == 0.0:
|
||||
top_p = 1.0 # Default top_p
|
||||
|
||||
# Initialize generation parameters for MLX-VLM
|
||||
generation_params = {
|
||||
max_tokens = getattr(request, 'Tokens', default_max_tokens) or default_max_tokens
|
||||
|
||||
temp = getattr(request, 'Temperature', 0.0) or 0.6
|
||||
top_p = getattr(request, 'TopP', 0.0) or 1.0
|
||||
min_p = getattr(request, 'MinP', 0.0) or 0.0
|
||||
top_k = getattr(request, 'TopK', 0) or 0
|
||||
|
||||
sampler_params = {
|
||||
'temp': temp,
|
||||
'top_p': top_p,
|
||||
'min_p': min_p,
|
||||
'top_k': top_k,
|
||||
}
|
||||
|
||||
# Add seed if specified
|
||||
|
||||
logits_params = {}
|
||||
repetition_penalty = getattr(request, 'RepetitionPenalty', 0.0) or 0.0
|
||||
if repetition_penalty and repetition_penalty != 1.0:
|
||||
logits_params['repetition_penalty'] = repetition_penalty
|
||||
presence_penalty = getattr(request, 'PresencePenalty', 0.0) or 0.0
|
||||
if presence_penalty:
|
||||
logits_params['presence_penalty'] = presence_penalty
|
||||
frequency_penalty = getattr(request, 'FrequencyPenalty', 0.0) or 0.0
|
||||
if frequency_penalty:
|
||||
logits_params['frequency_penalty'] = frequency_penalty
|
||||
|
||||
seed = getattr(request, 'Seed', 0)
|
||||
if seed != 0:
|
||||
mx.random.seed(seed)
|
||||
|
||||
# Override with options if available
|
||||
|
||||
if hasattr(self, 'options'):
|
||||
# Max tokens from options
|
||||
if 'max_tokens' in self.options:
|
||||
max_tokens = self.options['max_tokens']
|
||||
|
||||
# Generation parameters from options
|
||||
param_option_mapping = {
|
||||
'temp': 'temp',
|
||||
'temperature': 'temp', # alias
|
||||
'top_p': 'top_p',
|
||||
option_mapping = {
|
||||
'temp': 'temp', 'temperature': 'temp',
|
||||
'top_p': 'top_p', 'min_p': 'min_p', 'top_k': 'top_k',
|
||||
}
|
||||
|
||||
for option_key, param_key in param_option_mapping.items():
|
||||
for option_key, param_key in option_mapping.items():
|
||||
if option_key in self.options:
|
||||
generation_params[param_key] = self.options[option_key]
|
||||
|
||||
# Handle seed from options
|
||||
sampler_params[param_key] = self.options[option_key]
|
||||
for option_key in ('repetition_penalty', 'presence_penalty', 'frequency_penalty'):
|
||||
if option_key in self.options:
|
||||
logits_params[option_key] = self.options[option_key]
|
||||
if 'seed' in self.options:
|
||||
mx.random.seed(self.options['seed'])
|
||||
|
||||
return max_tokens, generation_params
|
||||
|
||||
stop_words = list(getattr(request, 'StopPrompts', []) or [])
|
||||
return max_tokens, sampler_params, logits_params, stop_words
|
||||
|
||||
def _finalize_output(self, request, generated_text, last_response):
|
||||
"""Split reasoning + tool calls out of generated_text and return the
|
||||
tuple consumed by Reply-builders."""
|
||||
content = generated_text
|
||||
reasoning_content = ""
|
||||
|
||||
if getattr(self, "has_thinking", False):
|
||||
reasoning_content, content = split_reasoning(content, self.think_start, self.think_end)
|
||||
|
||||
tool_calls_proto: List[backend_pb2.ToolCallDelta] = []
|
||||
if self.tool_module is not None:
|
||||
parsed_tools = None
|
||||
if request.Tools:
|
||||
try:
|
||||
parsed_tools = json.loads(request.Tools)
|
||||
except json.JSONDecodeError:
|
||||
parsed_tools = None
|
||||
calls, content = parse_tool_calls(content, self.tool_module, parsed_tools)
|
||||
for c in calls:
|
||||
tool_calls_proto.append(
|
||||
backend_pb2.ToolCallDelta(
|
||||
index=c["index"],
|
||||
id=c["id"],
|
||||
name=c["name"],
|
||||
arguments=c["arguments"],
|
||||
)
|
||||
)
|
||||
|
||||
prompt_tokens = int(getattr(last_response, "prompt_tokens", 0) or 0) if last_response else 0
|
||||
completion_tokens = int(getattr(last_response, "generation_tokens", 0) or 0) if last_response else 0
|
||||
|
||||
logprobs_bytes = b""
|
||||
if last_response is not None and int(getattr(request, "Logprobs", 0) or 0) > 0:
|
||||
try:
|
||||
lp = getattr(last_response, "logprobs", None)
|
||||
if lp is not None:
|
||||
token_id = int(getattr(last_response, "token", 0) or 0)
|
||||
tokenizer = (
|
||||
self.processor.tokenizer
|
||||
if hasattr(self.processor, "tokenizer")
|
||||
else self.processor
|
||||
)
|
||||
token_text = tokenizer.decode([token_id]) if token_id else ""
|
||||
top_logprob = float(lp[token_id]) if hasattr(lp, "__getitem__") else 0.0
|
||||
logprobs_bytes = json.dumps(
|
||||
{"content": [{"token": token_text, "logprob": top_logprob}]}
|
||||
).encode("utf-8")
|
||||
except Exception as e:
|
||||
print(f"[mlx-vlm] Logprobs extraction failed: {e}", file=sys.stderr)
|
||||
|
||||
return content, reasoning_content, tool_calls_proto, prompt_tokens, completion_tokens, logprobs_bytes
|
||||
|
||||
def _truncate_at_stop(self, text, stop_words):
|
||||
if not stop_words:
|
||||
return text
|
||||
earliest = len(text)
|
||||
for stop in stop_words:
|
||||
if not stop:
|
||||
continue
|
||||
idx = text.find(stop)
|
||||
if idx >= 0 and idx < earliest:
|
||||
earliest = idx
|
||||
return text[:earliest] if earliest < len(text) else text
|
||||
|
||||
def load_image_from_base64(self, image_data: str):
|
||||
"""
|
||||
|
||||
@@ -1,17 +1,19 @@
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
import subprocess
|
||||
import time
|
||||
|
||||
import grpc
|
||||
import backend_pb2
|
||||
import backend_pb2_grpc
|
||||
|
||||
import grpc
|
||||
|
||||
import unittest
|
||||
import subprocess
|
||||
import time
|
||||
import grpc
|
||||
import backend_pb2_grpc
|
||||
import backend_pb2
|
||||
# Make the shared helpers importable so we can unit-test them without a
|
||||
# running gRPC server.
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
|
||||
from python_utils import messages_to_dicts, parse_options
|
||||
from mlx_utils import parse_tool_calls, split_reasoning
|
||||
|
||||
class TestBackendServicer(unittest.TestCase):
|
||||
"""
|
||||
@@ -143,4 +145,55 @@ class TestBackendServicer(unittest.TestCase):
|
||||
print(err)
|
||||
self.fail("Embedding service failed")
|
||||
finally:
|
||||
self.tearDown()
|
||||
self.tearDown()
|
||||
|
||||
|
||||
class TestSharedHelpers(unittest.TestCase):
|
||||
"""Server-less unit tests for the helpers the mlx-vlm backend depends on."""
|
||||
|
||||
def test_parse_options_typed(self):
|
||||
opts = parse_options(["temperature:0.7", "max_tokens:128", "trust:true", "name:hello"])
|
||||
self.assertEqual(opts["temperature"], 0.7)
|
||||
self.assertEqual(opts["max_tokens"], 128)
|
||||
self.assertIs(opts["trust"], True)
|
||||
self.assertEqual(opts["name"], "hello")
|
||||
|
||||
def test_messages_to_dicts_roundtrip(self):
|
||||
msgs = [
|
||||
backend_pb2.Message(role="user", content="hi"),
|
||||
backend_pb2.Message(
|
||||
role="assistant",
|
||||
content="",
|
||||
tool_calls='[{"id":"call_1","type":"function","function":{"name":"f","arguments":"{}"}}]',
|
||||
),
|
||||
backend_pb2.Message(
|
||||
role="tool",
|
||||
content="42",
|
||||
tool_call_id="call_1",
|
||||
name="f",
|
||||
),
|
||||
]
|
||||
out = messages_to_dicts(msgs)
|
||||
self.assertEqual(out[0], {"role": "user", "content": "hi"})
|
||||
self.assertEqual(out[1]["tool_calls"][0]["function"]["name"], "f")
|
||||
self.assertEqual(out[2]["tool_call_id"], "call_1")
|
||||
|
||||
def test_split_reasoning(self):
|
||||
r, c = split_reasoning("<think>plan</think>final", "<think>", "</think>")
|
||||
self.assertEqual(r, "plan")
|
||||
self.assertEqual(c, "final")
|
||||
|
||||
def test_parse_tool_calls_with_shim(self):
|
||||
tm = types.SimpleNamespace(
|
||||
tool_call_start="<tool_call>",
|
||||
tool_call_end="</tool_call>",
|
||||
parse_tool_call=lambda body, tools: {"name": "get_weather", "arguments": {"location": body.strip()}},
|
||||
)
|
||||
calls, remaining = parse_tool_calls(
|
||||
"<tool_call>Paris</tool_call>",
|
||||
tm,
|
||||
tools=None,
|
||||
)
|
||||
self.assertEqual(len(calls), 1)
|
||||
self.assertEqual(calls[0]["name"], "get_weather")
|
||||
self.assertEqual(calls[0]["arguments"], '{"location": "Paris"}')
|
||||
@@ -2,11 +2,13 @@
|
||||
import asyncio
|
||||
from concurrent import futures
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import signal
|
||||
import sys
|
||||
import os
|
||||
import types
|
||||
from typing import List
|
||||
import time
|
||||
|
||||
import backend_pb2
|
||||
import backend_pb2_grpc
|
||||
@@ -15,13 +17,13 @@ import grpc
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
|
||||
from grpc_auth import get_auth_interceptors
|
||||
from python_utils import messages_to_dicts, parse_options
|
||||
from mlx_utils import parse_tool_calls, split_reasoning
|
||||
|
||||
from mlx_lm import load, generate, stream_generate
|
||||
from mlx_lm.sample_utils import make_sampler
|
||||
from mlx_lm import load, stream_generate
|
||||
from mlx_lm.sample_utils import make_logits_processors, make_sampler
|
||||
from mlx_lm.models.cache import make_prompt_cache, can_trim_prompt_cache, trim_prompt_cache
|
||||
import mlx.core as mx
|
||||
import base64
|
||||
import io
|
||||
|
||||
from mlx_cache import ThreadSafeLRUPromptCache
|
||||
|
||||
@@ -30,21 +32,6 @@ _ONE_DAY_IN_SECONDS = 60 * 60 * 24
|
||||
# If MAX_WORKERS are specified in the environment use it, otherwise default to 1
|
||||
MAX_WORKERS = int(os.environ.get('PYTHON_GRPC_MAX_WORKERS', '1'))
|
||||
|
||||
def is_float(s):
|
||||
"""Check if a string can be converted to float."""
|
||||
try:
|
||||
float(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
def is_int(s):
|
||||
"""Check if a string can be converted to int."""
|
||||
try:
|
||||
int(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
# Implement the BackendServicer class with the service methods
|
||||
class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
"""
|
||||
@@ -78,46 +65,27 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
try:
|
||||
print(f"Loading MLX model: {request.Model}", file=sys.stderr)
|
||||
print(f"Request: {request}", file=sys.stderr)
|
||||
|
||||
# Parse options like in the diffusers backend
|
||||
options = request.Options
|
||||
self.options = {}
|
||||
|
||||
# The options are a list of strings in this form optname:optvalue
|
||||
# We store all the options in a dict for later use
|
||||
for opt in options:
|
||||
if ":" not in opt:
|
||||
continue
|
||||
key, value = opt.split(":", 1) # Split only on first colon to handle values with colons
|
||||
|
||||
# Convert numeric values to appropriate types
|
||||
if is_float(value):
|
||||
value = float(value)
|
||||
elif is_int(value):
|
||||
value = int(value)
|
||||
elif value.lower() in ["true", "false"]:
|
||||
value = value.lower() == "true"
|
||||
|
||||
self.options[key] = value
|
||||
|
||||
|
||||
# Parse Options[] key:value strings into a typed dict (shared helper)
|
||||
self.options = parse_options(request.Options)
|
||||
print(f"Options: {self.options}", file=sys.stderr)
|
||||
|
||||
|
||||
# Build tokenizer config for MLX using options
|
||||
tokenizer_config = {}
|
||||
|
||||
|
||||
# Handle trust_remote_code from request or options
|
||||
if request.TrustRemoteCode or self.options.get("trust_remote_code", False):
|
||||
tokenizer_config["trust_remote_code"] = True
|
||||
|
||||
|
||||
# Handle EOS token from options
|
||||
if "eos_token" in self.options:
|
||||
tokenizer_config["eos_token"] = self.options["eos_token"]
|
||||
|
||||
|
||||
# Handle other tokenizer config options
|
||||
for key in ["pad_token", "bos_token", "unk_token", "sep_token", "cls_token", "mask_token"]:
|
||||
if key in self.options:
|
||||
tokenizer_config[key] = self.options[key]
|
||||
|
||||
|
||||
# Load model and tokenizer using MLX
|
||||
if tokenizer_config:
|
||||
print(f"Loading with tokenizer_config: {tokenizer_config}", file=sys.stderr)
|
||||
@@ -125,6 +93,21 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
else:
|
||||
self.model, self.tokenizer = load(request.Model)
|
||||
|
||||
# mlx_lm.load() returns a TokenizerWrapper that detects tool
|
||||
# calling and thinking markers from the chat template / vocab.
|
||||
# mlx-lm >= 0.30 also exposes a parser callable on the wrapper;
|
||||
# earlier versions don't (we fall back to json.loads inside
|
||||
# _tool_module_from_tokenizer below).
|
||||
has_tools = bool(getattr(self.tokenizer, "has_tool_calling", False))
|
||||
has_thinking = bool(getattr(self.tokenizer, "has_thinking", False))
|
||||
tcs = getattr(self.tokenizer, "tool_call_start", None)
|
||||
tce = getattr(self.tokenizer, "tool_call_end", None)
|
||||
print(
|
||||
f"MLX tokenizer capabilities: has_tool_calling={has_tools} "
|
||||
f"has_thinking={has_thinking} tool_call_start={tcs!r} tool_call_end={tce!r}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
# Initialize thread-safe LRU prompt cache for efficient generation
|
||||
max_cache_entries = self.options.get("max_cache_entries", 10)
|
||||
self.max_kv_size = self.options.get("max_kv_size", None)
|
||||
@@ -134,7 +117,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
can_trim_fn=can_trim_prompt_cache,
|
||||
trim_fn=trim_prompt_cache,
|
||||
)
|
||||
|
||||
|
||||
except Exception as err:
|
||||
print(f"Error loading MLX model {err=}, {type(err)=}", file=sys.stderr)
|
||||
return backend_pb2.Result(success=False, message=f"Error loading MLX model: {err}")
|
||||
@@ -172,30 +155,58 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
remaining_tokens = cache_key
|
||||
|
||||
# Build generation parameters using request attributes and options
|
||||
max_tokens, sampler_params = self._build_generation_params(request)
|
||||
max_tokens, sampler_params, logits_params, stop_words = self._build_generation_params(request)
|
||||
|
||||
print(f"Generating text with MLX - max_tokens: {max_tokens}, cache_hit: {len(remaining_tokens) < len(cache_key)}", file=sys.stderr)
|
||||
print(
|
||||
f"Generating text with MLX - max_tokens: {max_tokens}, "
|
||||
f"cache_hit: {len(remaining_tokens) < len(cache_key)}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
# Create sampler with parameters
|
||||
# Create sampler and optional logits processors (penalties)
|
||||
sampler = make_sampler(**sampler_params)
|
||||
logits_processors = make_logits_processors(**logits_params) if logits_params else None
|
||||
|
||||
# Use stream_generate to track generated tokens for cache key
|
||||
# Use stream_generate to collect text + track tokens for cache key
|
||||
generated_text = []
|
||||
last_response = None
|
||||
for response in stream_generate(
|
||||
self.model,
|
||||
self.tokenizer,
|
||||
prompt=remaining_tokens if remaining_tokens else cache_key,
|
||||
max_tokens=max_tokens,
|
||||
sampler=sampler,
|
||||
logits_processors=logits_processors,
|
||||
prompt_cache=prompt_cache,
|
||||
):
|
||||
generated_text.append(response.text)
|
||||
cache_key.append(response.token)
|
||||
last_response = response
|
||||
# Early stop on user-provided stop sequences
|
||||
if stop_words and any(s in "".join(generated_text) for s in stop_words):
|
||||
break
|
||||
|
||||
# Insert completed cache
|
||||
self.lru_cache.insert_cache(self.model_key, cache_key, prompt_cache)
|
||||
|
||||
return backend_pb2.Reply(message=bytes(''.join(generated_text), encoding='utf-8'))
|
||||
full_text = self._truncate_at_stop("".join(generated_text), stop_words)
|
||||
content, reasoning_content, tool_calls_proto, prompt_tokens, completion_tokens, logprobs_bytes = (
|
||||
self._finalize_output(request, full_text, last_response)
|
||||
)
|
||||
|
||||
return backend_pb2.Reply(
|
||||
message=bytes(content, encoding='utf-8'),
|
||||
prompt_tokens=prompt_tokens,
|
||||
tokens=completion_tokens,
|
||||
logprobs=logprobs_bytes,
|
||||
chat_deltas=[
|
||||
backend_pb2.ChatDelta(
|
||||
content=content,
|
||||
reasoning_content=reasoning_content,
|
||||
tool_calls=tool_calls_proto,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in MLX Predict: {e}", file=sys.stderr)
|
||||
@@ -206,7 +217,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
def Embedding(self, request, context):
|
||||
"""
|
||||
A gRPC method that calculates embeddings for a given sentence.
|
||||
|
||||
|
||||
Note: MLX-LM doesn't support embeddings directly. This method returns an error.
|
||||
|
||||
Args:
|
||||
@@ -221,6 +232,62 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
context.set_details("Embeddings are not supported in the MLX backend.")
|
||||
return backend_pb2.EmbeddingResult()
|
||||
|
||||
async def TokenizeString(self, request, context):
|
||||
"""Tokenize ``request.Prompt`` using the loaded model's tokenizer."""
|
||||
if not hasattr(self, "tokenizer") or self.tokenizer is None:
|
||||
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
|
||||
context.set_details("tokenizer not loaded")
|
||||
return backend_pb2.TokenizationResponse()
|
||||
try:
|
||||
tokens = self.tokenizer.encode(request.Prompt)
|
||||
if hasattr(tokens, "tolist"):
|
||||
tokens = tokens.tolist()
|
||||
tokens = list(tokens)
|
||||
return backend_pb2.TokenizationResponse(length=len(tokens), tokens=tokens)
|
||||
except Exception as e:
|
||||
context.set_code(grpc.StatusCode.INTERNAL)
|
||||
context.set_details(str(e))
|
||||
return backend_pb2.TokenizationResponse()
|
||||
|
||||
async def Free(self, request, context):
|
||||
"""Drop the loaded model, tokenizer and prompt cache.
|
||||
|
||||
Metal / CUDA memory is released via ``gc.collect()`` + the
|
||||
platform-specific cache clear hooks when available.
|
||||
"""
|
||||
try:
|
||||
if hasattr(self, "model"):
|
||||
del self.model
|
||||
if hasattr(self, "tokenizer"):
|
||||
del self.tokenizer
|
||||
if hasattr(self, "lru_cache") and self.lru_cache is not None:
|
||||
try:
|
||||
self.lru_cache.clear()
|
||||
except Exception:
|
||||
pass
|
||||
self.lru_cache = None
|
||||
gc.collect()
|
||||
# Metal: drop the cached allocator. mlx.clear_cache (mlx >= 0.30)
|
||||
# supersedes the now-deprecated mlx.metal.clear_cache.
|
||||
try:
|
||||
if hasattr(mx, "clear_cache"):
|
||||
mx.clear_cache()
|
||||
elif hasattr(mx, "metal") and hasattr(mx.metal, "clear_cache"):
|
||||
mx.metal.clear_cache()
|
||||
except Exception:
|
||||
pass
|
||||
# CUDA: release the torch cache if a CUDA-backed mlx variant
|
||||
# happens to be installed alongside torch (best-effort).
|
||||
try:
|
||||
import torch # type: ignore
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
except Exception:
|
||||
pass
|
||||
return backend_pb2.Result(success=True, message="MLX model freed")
|
||||
except Exception as e:
|
||||
return backend_pb2.Result(success=False, message=str(e))
|
||||
|
||||
async def PredictStream(self, request, context):
|
||||
"""
|
||||
Generates text based on the given prompt and sampling parameters, and streams the results using MLX.
|
||||
@@ -251,24 +318,64 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
remaining_tokens = cache_key
|
||||
|
||||
# Build generation parameters using request attributes and options
|
||||
max_tokens, sampler_params = self._build_generation_params(request, default_max_tokens=512)
|
||||
max_tokens, sampler_params, logits_params, stop_words = self._build_generation_params(
|
||||
request, default_max_tokens=512
|
||||
)
|
||||
|
||||
print(f"Streaming text with MLX - max_tokens: {max_tokens}, cache_hit: {len(remaining_tokens) < len(cache_key)}", file=sys.stderr)
|
||||
print(
|
||||
f"Streaming text with MLX - max_tokens: {max_tokens}, "
|
||||
f"cache_hit: {len(remaining_tokens) < len(cache_key)}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
# Create sampler with parameters
|
||||
# Create sampler and optional logits processors (penalties)
|
||||
sampler = make_sampler(**sampler_params)
|
||||
logits_processors = make_logits_processors(**logits_params) if logits_params else None
|
||||
|
||||
# Stream text generation using MLX with proper parameters
|
||||
accumulated = []
|
||||
last_response = None
|
||||
for response in stream_generate(
|
||||
self.model,
|
||||
self.tokenizer,
|
||||
prompt=remaining_tokens if remaining_tokens else cache_key,
|
||||
max_tokens=max_tokens,
|
||||
sampler=sampler,
|
||||
logits_processors=logits_processors,
|
||||
prompt_cache=prompt_cache,
|
||||
):
|
||||
cache_key.append(response.token)
|
||||
yield backend_pb2.Reply(message=bytes(response.text, encoding='utf-8'))
|
||||
accumulated.append(response.text)
|
||||
last_response = response
|
||||
# Emit a content delta. Structured reasoning / tool parsing
|
||||
# happens on the final chunk so we don't fragment the state
|
||||
# machine in v1.
|
||||
yield backend_pb2.Reply(
|
||||
message=bytes(response.text, encoding='utf-8'),
|
||||
chat_deltas=[backend_pb2.ChatDelta(content=response.text)],
|
||||
)
|
||||
# Early stop on user-provided stop sequences
|
||||
if stop_words and any(s in "".join(accumulated) for s in stop_words):
|
||||
break
|
||||
|
||||
# Final chunk: run reasoning + tool parsing on accumulated text
|
||||
# and emit the structured ChatDelta with token counts + logprobs.
|
||||
full_text = self._truncate_at_stop("".join(accumulated), stop_words)
|
||||
content, reasoning_content, tool_calls_proto, prompt_tokens, completion_tokens, logprobs_bytes = (
|
||||
self._finalize_output(request, full_text, last_response)
|
||||
)
|
||||
yield backend_pb2.Reply(
|
||||
message=b"",
|
||||
prompt_tokens=prompt_tokens,
|
||||
tokens=completion_tokens,
|
||||
logprobs=logprobs_bytes,
|
||||
chat_deltas=[
|
||||
backend_pb2.ChatDelta(
|
||||
content="",
|
||||
reasoning_content=reasoning_content,
|
||||
tool_calls=tool_calls_proto,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in MLX PredictStream: {e}", file=sys.stderr)
|
||||
@@ -294,21 +401,33 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
Returns:
|
||||
str: The prepared prompt.
|
||||
"""
|
||||
# If tokenizer template is enabled and messages are provided instead of prompt, apply the tokenizer template
|
||||
# If tokenizer template is enabled and messages are provided instead
|
||||
# of prompt, apply the tokenizer template (forwards tool definitions
|
||||
# and enable_thinking when the model supports them).
|
||||
if not request.Prompt and request.UseTokenizerTemplate and request.Messages:
|
||||
# Convert gRPC messages to the format expected by apply_chat_template
|
||||
messages = []
|
||||
for msg in request.Messages:
|
||||
messages.append({"role": msg.role, "content": msg.content})
|
||||
messages = messages_to_dicts(request.Messages)
|
||||
|
||||
prompt = self.tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True
|
||||
)
|
||||
return prompt
|
||||
else:
|
||||
return request.Prompt
|
||||
kwargs = {"tokenize": False, "add_generation_prompt": True}
|
||||
if request.Tools:
|
||||
try:
|
||||
kwargs["tools"] = json.loads(request.Tools)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
enable_thinking = request.Metadata.get("enable_thinking", "").lower()
|
||||
if enable_thinking == "true":
|
||||
kwargs["enable_thinking"] = True
|
||||
|
||||
try:
|
||||
return self.tokenizer.apply_chat_template(messages, **kwargs)
|
||||
except TypeError:
|
||||
# Fallback for tokenizers whose template doesn't accept
|
||||
# tools= or enable_thinking=.
|
||||
return self.tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
)
|
||||
return request.Prompt
|
||||
|
||||
def _get_tokens_from_prompt(self, prompt_text: str) -> List[int]:
|
||||
"""
|
||||
@@ -338,18 +457,19 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
default_max_tokens: Default max_tokens if not specified.
|
||||
|
||||
Returns:
|
||||
tuple: (max_tokens, sampler_params dict)
|
||||
tuple: (max_tokens, sampler_params dict, logits_processor_params dict,
|
||||
stop_words list)
|
||||
"""
|
||||
# Extract max_tokens
|
||||
max_tokens = getattr(request, 'Tokens', default_max_tokens)
|
||||
if max_tokens == 0:
|
||||
max_tokens = default_max_tokens
|
||||
|
||||
|
||||
# Extract sampler parameters from request attributes
|
||||
temp = getattr(request, 'Temperature', 0.0)
|
||||
if temp == 0.0:
|
||||
temp = 0.6 # Default temperature
|
||||
|
||||
|
||||
top_p = getattr(request, 'TopP', 0.0)
|
||||
if top_p == 0.0:
|
||||
top_p = 1.0 # Default top_p
|
||||
@@ -369,18 +489,31 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
'xtc_threshold': 0.0,
|
||||
'xtc_probability': 0.0,
|
||||
}
|
||||
|
||||
|
||||
# Logits processor parameters — only set fields the request actually
|
||||
# provides so we can feed them unconditionally to make_logits_processors.
|
||||
logits_params = {}
|
||||
repetition_penalty = getattr(request, 'RepetitionPenalty', 0.0) or 0.0
|
||||
if repetition_penalty and repetition_penalty != 1.0:
|
||||
logits_params['repetition_penalty'] = repetition_penalty
|
||||
presence_penalty = getattr(request, 'PresencePenalty', 0.0) or 0.0
|
||||
if presence_penalty:
|
||||
logits_params['presence_penalty'] = presence_penalty
|
||||
frequency_penalty = getattr(request, 'FrequencyPenalty', 0.0) or 0.0
|
||||
if frequency_penalty:
|
||||
logits_params['frequency_penalty'] = frequency_penalty
|
||||
|
||||
# Add seed if specified
|
||||
seed = getattr(request, 'Seed', 0)
|
||||
if seed != 0:
|
||||
mx.random.seed(seed)
|
||||
|
||||
|
||||
# Override with options if available
|
||||
if hasattr(self, 'options'):
|
||||
# Max tokens from options
|
||||
if 'max_tokens' in self.options:
|
||||
max_tokens = self.options['max_tokens']
|
||||
|
||||
|
||||
# Sampler parameters from options
|
||||
sampler_option_mapping = {
|
||||
'temp': 'temp',
|
||||
@@ -391,32 +524,142 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
'xtc_threshold': 'xtc_threshold',
|
||||
'xtc_probability': 'xtc_probability',
|
||||
}
|
||||
|
||||
|
||||
for option_key, param_key in sampler_option_mapping.items():
|
||||
if option_key in self.options:
|
||||
sampler_params[param_key] = self.options[option_key]
|
||||
|
||||
|
||||
# Logits processor overrides
|
||||
for option_key in ('repetition_penalty', 'presence_penalty', 'frequency_penalty'):
|
||||
if option_key in self.options:
|
||||
logits_params[option_key] = self.options[option_key]
|
||||
|
||||
# Handle seed from options
|
||||
if 'seed' in self.options:
|
||||
mx.random.seed(self.options['seed'])
|
||||
|
||||
|
||||
# Special tokens for XTC sampling (if tokenizer has eos_token_ids)
|
||||
xtc_special_tokens = []
|
||||
if hasattr(self.tokenizer, 'eos_token_ids') and self.tokenizer.eos_token_ids:
|
||||
xtc_special_tokens = list(self.tokenizer.eos_token_ids)
|
||||
elif hasattr(self.tokenizer, 'eos_token_id') and self.tokenizer.eos_token_id is not None:
|
||||
xtc_special_tokens = [self.tokenizer.eos_token_id]
|
||||
|
||||
|
||||
# Add newline token if available
|
||||
try:
|
||||
newline_tokens = self.tokenizer.encode("\n")
|
||||
xtc_special_tokens.extend(newline_tokens)
|
||||
except:
|
||||
except Exception:
|
||||
pass # Skip if encoding fails
|
||||
|
||||
|
||||
sampler_params['xtc_special_tokens'] = xtc_special_tokens
|
||||
|
||||
return max_tokens, sampler_params
|
||||
|
||||
# Stop sequences are applied post-decode (mlx-lm doesn't have a
|
||||
# built-in stop-sequence sampler param). Preserve the list here.
|
||||
stop_words = list(getattr(request, 'StopPrompts', []) or [])
|
||||
|
||||
return max_tokens, sampler_params, logits_params, stop_words
|
||||
|
||||
def _tool_module_from_tokenizer(self):
|
||||
"""Build a duck-typed tool module from the TokenizerWrapper.
|
||||
|
||||
On mlx-lm >= 0.30 the wrapper exposes a ``tool_parser`` callable
|
||||
that's been resolved from the model's chat template. On older
|
||||
releases (e.g. 0.29.x) the wrapper only carries the start/end
|
||||
markers — fall back to ``json.loads`` of the body, which matches
|
||||
what ``mlx_lm.tool_parsers.json_tools.parse_tool_call`` does on
|
||||
HEAD and covers the only format 0.29 detects (``<tool_call>``).
|
||||
"""
|
||||
start = getattr(self.tokenizer, "tool_call_start", None)
|
||||
end = getattr(self.tokenizer, "tool_call_end", None)
|
||||
if not start:
|
||||
return None
|
||||
parse_fn = getattr(self.tokenizer, "tool_parser", None)
|
||||
if parse_fn is None:
|
||||
def parse_fn(body, tools): # noqa: E306 — local fallback
|
||||
return json.loads(body.strip())
|
||||
return types.SimpleNamespace(
|
||||
tool_call_start=start,
|
||||
tool_call_end=end or "",
|
||||
parse_tool_call=parse_fn,
|
||||
)
|
||||
|
||||
def _finalize_output(self, request, generated_text, last_response):
|
||||
"""Build a ChatDelta + token counts + logprobs from accumulated output.
|
||||
|
||||
Returns ``(content, reasoning_content, tool_calls_proto,
|
||||
prompt_token_count, completion_token_count, logprobs_bytes)``.
|
||||
"""
|
||||
content = generated_text
|
||||
reasoning_content = ""
|
||||
|
||||
if getattr(self.tokenizer, "has_thinking", False):
|
||||
think_start = getattr(self.tokenizer, "think_start", "") or ""
|
||||
think_end = getattr(self.tokenizer, "think_end", "") or ""
|
||||
reasoning_content, content = split_reasoning(content, think_start, think_end)
|
||||
|
||||
tool_calls_proto: List[backend_pb2.ToolCallDelta] = []
|
||||
tool_module = None
|
||||
if getattr(self.tokenizer, "has_tool_calling", False):
|
||||
tool_module = self._tool_module_from_tokenizer()
|
||||
if tool_module is not None:
|
||||
parsed_tools = None
|
||||
if request.Tools:
|
||||
try:
|
||||
parsed_tools = json.loads(request.Tools)
|
||||
except json.JSONDecodeError:
|
||||
parsed_tools = None
|
||||
calls, content = parse_tool_calls(content, tool_module, parsed_tools)
|
||||
for c in calls:
|
||||
tool_calls_proto.append(
|
||||
backend_pb2.ToolCallDelta(
|
||||
index=c["index"],
|
||||
id=c["id"],
|
||||
name=c["name"],
|
||||
arguments=c["arguments"],
|
||||
)
|
||||
)
|
||||
|
||||
prompt_token_count = int(getattr(last_response, "prompt_tokens", 0) or 0) if last_response else 0
|
||||
completion_token_count = int(getattr(last_response, "generation_tokens", 0) or 0) if last_response else 0
|
||||
|
||||
logprobs_bytes = b""
|
||||
# Logprobs extraction — only when the request asked for them.
|
||||
if last_response is not None and int(getattr(request, "Logprobs", 0) or 0) > 0:
|
||||
try:
|
||||
lp = getattr(last_response, "logprobs", None)
|
||||
if lp is not None:
|
||||
# GenerationResponse.logprobs on the last chunk is the
|
||||
# logprob distribution of the final token. Without a
|
||||
# per-token history we at minimum surface the last token's
|
||||
# top-1 logprob so clients get a non-empty field.
|
||||
token_id = int(getattr(last_response, "token", 0) or 0)
|
||||
token_text = self.tokenizer.decode([token_id]) if token_id else ""
|
||||
top_logprob = float(lp[token_id]) if hasattr(lp, "__getitem__") else 0.0
|
||||
logprobs_bytes = json.dumps(
|
||||
{
|
||||
"content": [
|
||||
{"token": token_text, "logprob": top_logprob}
|
||||
]
|
||||
}
|
||||
).encode("utf-8")
|
||||
except Exception as e:
|
||||
print(f"[mlx] Logprobs extraction failed: {e}", file=sys.stderr)
|
||||
|
||||
return content, reasoning_content, tool_calls_proto, prompt_token_count, completion_token_count, logprobs_bytes
|
||||
|
||||
def _truncate_at_stop(self, text, stop_words):
|
||||
"""Truncate ``text`` at the first occurrence of any stop sequence."""
|
||||
if not stop_words:
|
||||
return text
|
||||
earliest = len(text)
|
||||
for stop in stop_words:
|
||||
if not stop:
|
||||
continue
|
||||
idx = text.find(stop)
|
||||
if idx >= 0 and idx < earliest:
|
||||
earliest = idx
|
||||
return text[:earliest] if earliest < len(text) else text
|
||||
|
||||
async def serve(address):
|
||||
# Start asyncio gRPC server
|
||||
|
||||
@@ -1,11 +1,20 @@
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
import subprocess
|
||||
import time
|
||||
import types
|
||||
|
||||
import grpc
|
||||
import backend_pb2
|
||||
import backend_pb2_grpc
|
||||
|
||||
# Make the shared helpers importable so we can unit-test them without a
|
||||
# running gRPC server.
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
|
||||
from python_utils import messages_to_dicts, parse_options
|
||||
from mlx_utils import parse_tool_calls, split_reasoning
|
||||
|
||||
class TestBackendServicer(unittest.TestCase):
|
||||
"""
|
||||
TestBackendServicer is the class that tests the gRPC service.
|
||||
@@ -231,4 +240,104 @@ class TestBackendServicer(unittest.TestCase):
|
||||
self.tearDown()
|
||||
|
||||
|
||||
def test_tokenize_string(self):
|
||||
"""TokenizeString should return a non-empty token list for a known prompt."""
|
||||
try:
|
||||
self.setUp()
|
||||
with grpc.insecure_channel("localhost:50051") as channel:
|
||||
stub = backend_pb2_grpc.BackendStub(channel)
|
||||
response = stub.LoadModel(
|
||||
backend_pb2.ModelOptions(Model="mlx-community/Llama-3.2-1B-Instruct-4bit")
|
||||
)
|
||||
self.assertTrue(response.success)
|
||||
resp = stub.TokenizeString(backend_pb2.PredictOptions(Prompt="Hello, world"))
|
||||
self.assertGreater(resp.length, 0)
|
||||
self.assertEqual(len(list(resp.tokens)), resp.length)
|
||||
except Exception as err:
|
||||
print(err)
|
||||
self.fail("TokenizeString service failed")
|
||||
finally:
|
||||
self.tearDown()
|
||||
|
||||
def test_free(self):
|
||||
"""Free should release the model and not crash on subsequent calls."""
|
||||
try:
|
||||
self.setUp()
|
||||
with grpc.insecure_channel("localhost:50051") as channel:
|
||||
stub = backend_pb2_grpc.BackendStub(channel)
|
||||
response = stub.LoadModel(
|
||||
backend_pb2.ModelOptions(Model="mlx-community/Llama-3.2-1B-Instruct-4bit")
|
||||
)
|
||||
self.assertTrue(response.success)
|
||||
free_resp = stub.Free(backend_pb2.HealthMessage())
|
||||
self.assertTrue(free_resp.success)
|
||||
except Exception as err:
|
||||
print(err)
|
||||
self.fail("Free service failed")
|
||||
finally:
|
||||
self.tearDown()
|
||||
|
||||
|
||||
class TestSharedHelpers(unittest.TestCase):
|
||||
"""Server-less unit tests for the helpers the mlx backend depends on."""
|
||||
|
||||
def test_parse_options_typed(self):
|
||||
opts = parse_options(["temperature:0.7", "max_tokens:128", "trust:true", "name:hello", "no_colon_skipped"])
|
||||
self.assertEqual(opts["temperature"], 0.7)
|
||||
self.assertEqual(opts["max_tokens"], 128)
|
||||
self.assertIs(opts["trust"], True)
|
||||
self.assertEqual(opts["name"], "hello")
|
||||
self.assertNotIn("no_colon_skipped", opts)
|
||||
|
||||
def test_messages_to_dicts_roundtrip(self):
|
||||
# Build proto Message objects (via backend_pb2 to match real gRPC)
|
||||
msgs = [
|
||||
backend_pb2.Message(role="user", content="hi"),
|
||||
backend_pb2.Message(
|
||||
role="assistant",
|
||||
content="",
|
||||
tool_calls='[{"id":"call_1","type":"function","function":{"name":"f","arguments":"{}"}}]',
|
||||
),
|
||||
backend_pb2.Message(
|
||||
role="tool",
|
||||
content="42",
|
||||
tool_call_id="call_1",
|
||||
name="f",
|
||||
),
|
||||
]
|
||||
out = messages_to_dicts(msgs)
|
||||
self.assertEqual(out[0], {"role": "user", "content": "hi"})
|
||||
self.assertEqual(out[1]["role"], "assistant")
|
||||
self.assertEqual(out[1]["tool_calls"][0]["function"]["name"], "f")
|
||||
self.assertEqual(out[2]["tool_call_id"], "call_1")
|
||||
self.assertEqual(out[2]["name"], "f")
|
||||
|
||||
def test_split_reasoning(self):
|
||||
r, c = split_reasoning("<think>step 1\nstep 2</think>The answer is 42.", "<think>", "</think>")
|
||||
self.assertEqual(r, "step 1\nstep 2")
|
||||
self.assertEqual(c, "The answer is 42.")
|
||||
|
||||
def test_split_reasoning_no_marker(self):
|
||||
r, c = split_reasoning("just text", "<think>", "</think>")
|
||||
self.assertEqual(r, "")
|
||||
self.assertEqual(c, "just text")
|
||||
|
||||
def test_parse_tool_calls_with_shim(self):
|
||||
tm = types.SimpleNamespace(
|
||||
tool_call_start="<tool_call>",
|
||||
tool_call_end="</tool_call>",
|
||||
parse_tool_call=lambda body, tools: {"name": "get_weather", "arguments": {"location": body.strip()}},
|
||||
)
|
||||
calls, remaining = parse_tool_calls(
|
||||
"Sure: <tool_call>Paris</tool_call>",
|
||||
tm,
|
||||
tools=None,
|
||||
)
|
||||
self.assertEqual(len(calls), 1)
|
||||
self.assertEqual(calls[0]["name"], "get_weather")
|
||||
self.assertEqual(calls[0]["arguments"], '{"location": "Paris"}')
|
||||
self.assertEqual(calls[0]["index"], 0)
|
||||
self.assertNotIn("<tool_call>", remaining)
|
||||
|
||||
|
||||
# Unit tests for ThreadSafeLRUPromptCache are in test_mlx_cache.py
|
||||
17
backend/python/sglang/Makefile
Normal file
17
backend/python/sglang/Makefile
Normal file
@@ -0,0 +1,17 @@
|
||||
.PHONY: sglang
|
||||
sglang:
|
||||
bash install.sh
|
||||
|
||||
.PHONY: run
|
||||
run: sglang
|
||||
@echo "Running sglang..."
|
||||
bash run.sh
|
||||
@echo "sglang run."
|
||||
|
||||
.PHONY: protogen-clean
|
||||
protogen-clean:
|
||||
$(RM) backend_pb2_grpc.py backend_pb2.py
|
||||
|
||||
.PHONY: clean
|
||||
clean: protogen-clean
|
||||
rm -rf venv __pycache__
|
||||
502
backend/python/sglang/backend.py
Normal file
502
backend/python/sglang/backend.py
Normal file
@@ -0,0 +1,502 @@
|
||||
#!/usr/bin/env python3
|
||||
"""LocalAI gRPC backend for sglang.
|
||||
|
||||
Wraps sglang's async Engine API behind the Backend gRPC contract defined
|
||||
in backend.proto. Mirrors the structure of backend/python/vllm/backend.py
|
||||
so that the two backends stay behavior-equivalent at the protocol level.
|
||||
|
||||
The streaming path applies sglang's per-request FunctionCallParser and
|
||||
ReasoningParser so tool_calls and reasoning_content are emitted
|
||||
incrementally inside ChatDelta, which is a capability sglang exposes
|
||||
natively and vLLM does not.
|
||||
"""
|
||||
import asyncio
|
||||
from concurrent import futures
|
||||
import argparse
|
||||
import signal
|
||||
import sys
|
||||
import os
|
||||
import json
|
||||
import gc
|
||||
import uuid
|
||||
import base64
|
||||
import io
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from PIL import Image
|
||||
|
||||
import backend_pb2
|
||||
import backend_pb2_grpc
|
||||
|
||||
import grpc
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
|
||||
from grpc_auth import get_auth_interceptors
|
||||
|
||||
# sglang imports. Engine is the stable public entry point; parser modules
|
||||
# are wrapped in try/except so older / leaner installs that omit them
|
||||
# still load the backend for plain text generation.
|
||||
from sglang.srt.entrypoints.engine import Engine
|
||||
|
||||
try:
|
||||
from sglang.srt.function_call.function_call_parser import FunctionCallParser
|
||||
# sglang's FunctionCallParser expects a list of pydantic Tool objects
|
||||
# (protocol.Tool with .function.name), not plain dicts. Wrap at the
|
||||
# request boundary to match.
|
||||
from sglang.srt.entrypoints.openai.protocol import Tool as SglTool
|
||||
HAS_TOOL_PARSERS = True
|
||||
except Exception:
|
||||
FunctionCallParser = None # type: ignore
|
||||
SglTool = None # type: ignore
|
||||
HAS_TOOL_PARSERS = False
|
||||
|
||||
try:
|
||||
from sglang.srt.parser.reasoning_parser import ReasoningParser
|
||||
HAS_REASONING_PARSERS = True
|
||||
except Exception:
|
||||
ReasoningParser = None # type: ignore
|
||||
HAS_REASONING_PARSERS = False
|
||||
|
||||
try:
|
||||
from transformers import AutoTokenizer
|
||||
HAS_TRANSFORMERS = True
|
||||
except Exception:
|
||||
AutoTokenizer = None # type: ignore
|
||||
HAS_TRANSFORMERS = False
|
||||
|
||||
|
||||
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
|
||||
MAX_WORKERS = int(os.environ.get('PYTHON_GRPC_MAX_WORKERS', '1'))
|
||||
|
||||
|
||||
class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
"""gRPC servicer implementing the Backend service for sglang."""
|
||||
|
||||
def _parse_options(self, options_list) -> Dict[str, str]:
|
||||
opts: Dict[str, str] = {}
|
||||
for opt in options_list:
|
||||
if ":" not in opt:
|
||||
continue
|
||||
key, value = opt.split(":", 1)
|
||||
opts[key.strip()] = value.strip()
|
||||
return opts
|
||||
|
||||
def _messages_to_dicts(self, messages) -> List[dict]:
|
||||
result: List[dict] = []
|
||||
for msg in messages:
|
||||
d = {"role": msg.role, "content": msg.content or ""}
|
||||
if msg.name:
|
||||
d["name"] = msg.name
|
||||
if msg.tool_call_id:
|
||||
d["tool_call_id"] = msg.tool_call_id
|
||||
if msg.reasoning_content:
|
||||
d["reasoning_content"] = msg.reasoning_content
|
||||
if msg.tool_calls:
|
||||
try:
|
||||
d["tool_calls"] = json.loads(msg.tool_calls)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
result.append(d)
|
||||
return result
|
||||
|
||||
def Health(self, request, context):
|
||||
return backend_pb2.Reply(message=bytes("OK", 'utf-8'))
|
||||
|
||||
async def LoadModel(self, request, context):
|
||||
engine_kwargs = {"model_path": request.Model}
|
||||
|
||||
if request.Quantization:
|
||||
engine_kwargs["quantization"] = request.Quantization
|
||||
if request.LoadFormat:
|
||||
engine_kwargs["load_format"] = request.LoadFormat
|
||||
if request.GPUMemoryUtilization:
|
||||
engine_kwargs["mem_fraction_static"] = float(request.GPUMemoryUtilization)
|
||||
if request.TrustRemoteCode:
|
||||
engine_kwargs["trust_remote_code"] = True
|
||||
if request.EnforceEager:
|
||||
engine_kwargs["disable_cuda_graph"] = True
|
||||
if request.TensorParallelSize:
|
||||
engine_kwargs["tp_size"] = int(request.TensorParallelSize)
|
||||
if request.MaxModelLen:
|
||||
engine_kwargs["context_length"] = int(request.MaxModelLen)
|
||||
if request.DType:
|
||||
engine_kwargs["dtype"] = request.DType
|
||||
|
||||
opts = self._parse_options(request.Options)
|
||||
|
||||
# Cache parser names — actual parser instances are created per
|
||||
# request because sglang's parsers are stateful.
|
||||
self.tool_parser_name: Optional[str] = opts.get("tool_parser") or None
|
||||
self.reasoning_parser_name: Optional[str] = opts.get("reasoning_parser") or None
|
||||
|
||||
# Also hand the parser names to sglang's engine so its HTTP/OAI
|
||||
# paths work identically if someone hits the engine directly.
|
||||
if self.tool_parser_name:
|
||||
engine_kwargs["tool_call_parser"] = self.tool_parser_name
|
||||
if self.reasoning_parser_name:
|
||||
engine_kwargs["reasoning_parser"] = self.reasoning_parser_name
|
||||
|
||||
try:
|
||||
self.llm = Engine(**engine_kwargs)
|
||||
except Exception as err:
|
||||
print(f"sglang Engine init failed: {err!r}", file=sys.stderr)
|
||||
return backend_pb2.Result(success=False, message=f"{err!r}")
|
||||
|
||||
# sglang does not expose a uniform get_tokenizer() off Engine.
|
||||
# Use transformers directly — same path sglang uses internally.
|
||||
self.tokenizer = None
|
||||
if HAS_TRANSFORMERS:
|
||||
try:
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(
|
||||
request.Model,
|
||||
trust_remote_code=bool(request.TrustRemoteCode),
|
||||
)
|
||||
except Exception as err:
|
||||
print(f"AutoTokenizer load failed (non-fatal): {err!r}", file=sys.stderr)
|
||||
|
||||
print("Model loaded successfully", file=sys.stderr)
|
||||
return backend_pb2.Result(message="Model loaded successfully", success=True)
|
||||
|
||||
async def Predict(self, request, context):
|
||||
gen = self._predict(request, context, streaming=False)
|
||||
res = await gen.__anext__()
|
||||
return res
|
||||
|
||||
async def PredictStream(self, request, context):
|
||||
iterations = self._predict(request, context, streaming=True)
|
||||
try:
|
||||
async for iteration in iterations:
|
||||
yield iteration
|
||||
finally:
|
||||
try:
|
||||
await iterations.aclose()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def TokenizeString(self, request, context):
|
||||
if not getattr(self, "tokenizer", None):
|
||||
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
|
||||
context.set_details("tokenizer not loaded")
|
||||
return backend_pb2.TokenizationResponse()
|
||||
try:
|
||||
tokens = self.tokenizer.encode(request.Prompt)
|
||||
return backend_pb2.TokenizationResponse(length=len(tokens), tokens=tokens)
|
||||
except Exception as e:
|
||||
context.set_code(grpc.StatusCode.INTERNAL)
|
||||
context.set_details(str(e))
|
||||
return backend_pb2.TokenizationResponse()
|
||||
|
||||
async def Free(self, request, context):
|
||||
try:
|
||||
if hasattr(self, "llm"):
|
||||
try:
|
||||
self.llm.shutdown()
|
||||
except Exception:
|
||||
pass
|
||||
del self.llm
|
||||
if hasattr(self, "tokenizer"):
|
||||
del self.tokenizer
|
||||
self.tool_parser_name = None
|
||||
self.reasoning_parser_name = None
|
||||
gc.collect()
|
||||
try:
|
||||
import torch
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
except ImportError:
|
||||
pass
|
||||
return backend_pb2.Result(success=True, message="Model freed")
|
||||
except Exception as e:
|
||||
return backend_pb2.Result(success=False, message=str(e))
|
||||
|
||||
def _build_sampling_params(self, request) -> dict:
|
||||
sampling_params: dict = {"temperature": 0.7, "max_new_tokens": 200}
|
||||
mapping = {
|
||||
"N": "n",
|
||||
"PresencePenalty": "presence_penalty",
|
||||
"FrequencyPenalty": "frequency_penalty",
|
||||
"RepetitionPenalty": "repetition_penalty",
|
||||
"Temperature": "temperature",
|
||||
"TopP": "top_p",
|
||||
"TopK": "top_k",
|
||||
"MinP": "min_p",
|
||||
"Seed": "seed",
|
||||
"StopPrompts": "stop",
|
||||
"StopTokenIds": "stop_token_ids",
|
||||
"IgnoreEOS": "ignore_eos",
|
||||
"Tokens": "max_new_tokens",
|
||||
"MinTokens": "min_new_tokens",
|
||||
"SkipSpecialTokens": "skip_special_tokens",
|
||||
}
|
||||
for proto_field, sgl_key in mapping.items():
|
||||
if not hasattr(request, proto_field):
|
||||
continue
|
||||
value = getattr(request, proto_field)
|
||||
if value in (None, 0, 0.0, [], False, ""):
|
||||
continue
|
||||
# repeated fields come back as RepeatedScalarContainer — convert
|
||||
if hasattr(value, "__iter__") and not isinstance(value, (str, bytes)):
|
||||
value = list(value)
|
||||
if not value:
|
||||
continue
|
||||
sampling_params[sgl_key] = value
|
||||
|
||||
# Grammar → JSON schema or EBNF structured decoding.
|
||||
if getattr(request, "Grammar", ""):
|
||||
grammar = request.Grammar
|
||||
try:
|
||||
json.loads(grammar)
|
||||
sampling_params["json_schema"] = grammar
|
||||
except json.JSONDecodeError:
|
||||
sampling_params["ebnf"] = grammar
|
||||
|
||||
return sampling_params
|
||||
|
||||
def _build_prompt(self, request) -> str:
|
||||
prompt = request.Prompt
|
||||
if prompt or not request.UseTokenizerTemplate or not request.Messages:
|
||||
return prompt
|
||||
|
||||
if self.tokenizer is None:
|
||||
print(
|
||||
"UseTokenizerTemplate requested but tokenizer not loaded; "
|
||||
"falling back to naive concatenation",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return "\n".join(m.content or "" for m in request.Messages)
|
||||
|
||||
messages_dicts = self._messages_to_dicts(request.Messages)
|
||||
template_kwargs: dict = {"tokenize": False, "add_generation_prompt": True}
|
||||
if request.Tools:
|
||||
try:
|
||||
template_kwargs["tools"] = json.loads(request.Tools)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
if request.Metadata.get("enable_thinking", "").lower() == "true":
|
||||
template_kwargs["enable_thinking"] = True
|
||||
|
||||
try:
|
||||
return self.tokenizer.apply_chat_template(messages_dicts, **template_kwargs)
|
||||
except TypeError:
|
||||
return self.tokenizer.apply_chat_template(
|
||||
messages_dicts, tokenize=False, add_generation_prompt=True,
|
||||
)
|
||||
|
||||
def _make_parsers(self, request):
|
||||
"""Construct fresh per-request parser instances (stateful)."""
|
||||
tool_parser = None
|
||||
reasoning_parser = None
|
||||
|
||||
if HAS_TOOL_PARSERS and self.tool_parser_name and request.Tools:
|
||||
try:
|
||||
tools_raw = json.loads(request.Tools)
|
||||
tools = [SglTool.model_validate(t) for t in tools_raw] if SglTool else tools_raw
|
||||
tool_parser = FunctionCallParser(
|
||||
tools=tools, tool_call_parser=self.tool_parser_name,
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"FunctionCallParser init failed: {e!r}", file=sys.stderr)
|
||||
|
||||
if HAS_REASONING_PARSERS and self.reasoning_parser_name:
|
||||
try:
|
||||
reasoning_parser = ReasoningParser(
|
||||
model_type=self.reasoning_parser_name,
|
||||
stream_reasoning=True,
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"ReasoningParser init failed: {e!r}", file=sys.stderr)
|
||||
|
||||
return tool_parser, reasoning_parser
|
||||
|
||||
async def _predict(self, request, context, streaming: bool = False):
|
||||
sampling_params = self._build_sampling_params(request)
|
||||
prompt = self._build_prompt(request)
|
||||
|
||||
tool_parser, reasoning_parser = self._make_parsers(request)
|
||||
|
||||
image_data = list(request.Images) if request.Images else None
|
||||
video_data = list(request.Videos) if request.Videos else None
|
||||
|
||||
# Kick off streaming generation. We always use stream=True so the
|
||||
# non-stream path still gets parser coverage on the final text.
|
||||
try:
|
||||
iterator = await self.llm.async_generate(
|
||||
prompt=prompt,
|
||||
sampling_params=sampling_params,
|
||||
image_data=image_data,
|
||||
video_data=video_data,
|
||||
stream=True,
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"sglang async_generate failed: {e!r}", file=sys.stderr)
|
||||
yield backend_pb2.Reply(message=bytes(f"error: {e!r}", "utf-8"))
|
||||
return
|
||||
|
||||
generated_text = ""
|
||||
last_chunk: Optional[dict] = None
|
||||
# Track tool call ids once per (request, tool_index) to match the
|
||||
# OpenAI streaming contract (id sent on first chunk for that tool).
|
||||
tool_ids_seen: Dict[int, str] = {}
|
||||
|
||||
try:
|
||||
async for chunk in iterator:
|
||||
last_chunk = chunk
|
||||
cumulative = chunk.get("text", "") if isinstance(chunk, dict) else ""
|
||||
delta_text = cumulative[len(generated_text):] if cumulative.startswith(generated_text) else cumulative
|
||||
generated_text = cumulative
|
||||
if not delta_text:
|
||||
continue
|
||||
|
||||
reasoning_delta = ""
|
||||
content_delta = delta_text
|
||||
|
||||
if reasoning_parser is not None:
|
||||
try:
|
||||
r, n = reasoning_parser.parse_stream_chunk(delta_text)
|
||||
reasoning_delta = r or ""
|
||||
content_delta = n or ""
|
||||
except Exception as e:
|
||||
print(f"reasoning_parser.parse_stream_chunk: {e!r}", file=sys.stderr)
|
||||
|
||||
tool_call_deltas: List[backend_pb2.ToolCallDelta] = []
|
||||
if tool_parser is not None and content_delta:
|
||||
try:
|
||||
normal_text, calls = tool_parser.parse_stream_chunk(content_delta)
|
||||
content_delta = normal_text or ""
|
||||
for tc in calls:
|
||||
idx = int(getattr(tc, "tool_index", 0) or 0)
|
||||
tc_id = tool_ids_seen.get(idx)
|
||||
if tc_id is None:
|
||||
tc_id = f"call_{uuid.uuid4().hex[:24]}"
|
||||
tool_ids_seen[idx] = tc_id
|
||||
tool_call_deltas.append(backend_pb2.ToolCallDelta(
|
||||
index=idx,
|
||||
id=tc_id,
|
||||
name=getattr(tc, "name", "") or "",
|
||||
arguments=getattr(tc, "parameters", "") or "",
|
||||
))
|
||||
except Exception as e:
|
||||
print(f"tool_parser.parse_stream_chunk: {e!r}", file=sys.stderr)
|
||||
|
||||
if streaming and (content_delta or reasoning_delta or tool_call_deltas):
|
||||
yield backend_pb2.Reply(
|
||||
message=bytes(content_delta, "utf-8"),
|
||||
chat_deltas=[backend_pb2.ChatDelta(
|
||||
content=content_delta,
|
||||
reasoning_content=reasoning_delta,
|
||||
tool_calls=tool_call_deltas,
|
||||
)],
|
||||
)
|
||||
finally:
|
||||
try:
|
||||
await iterator.aclose()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Extract token counts from the final chunk's meta_info.
|
||||
meta = {}
|
||||
if isinstance(last_chunk, dict):
|
||||
meta = last_chunk.get("meta_info") or {}
|
||||
prompt_tokens = int(meta.get("prompt_tokens", 0) or 0)
|
||||
completion_tokens = int(meta.get("completion_tokens", 0) or 0)
|
||||
|
||||
# Non-streaming path: re-parse the full text with fresh parsers
|
||||
# so we return a clean, complete ChatDelta. Streaming parsers
|
||||
# used above have accumulated state we don't want to reuse.
|
||||
final_content = generated_text
|
||||
final_reasoning = ""
|
||||
final_tool_calls: List[backend_pb2.ToolCallDelta] = []
|
||||
|
||||
if not streaming:
|
||||
final_reasoning_parser = None
|
||||
if HAS_REASONING_PARSERS and self.reasoning_parser_name:
|
||||
try:
|
||||
final_reasoning_parser = ReasoningParser(
|
||||
model_type=self.reasoning_parser_name,
|
||||
stream_reasoning=False,
|
||||
)
|
||||
except Exception:
|
||||
final_reasoning_parser = None
|
||||
|
||||
if final_reasoning_parser is not None:
|
||||
try:
|
||||
r, n = final_reasoning_parser.parse_non_stream(generated_text)
|
||||
final_reasoning = r or ""
|
||||
final_content = n if n is not None else generated_text
|
||||
except Exception as e:
|
||||
print(f"reasoning_parser.parse_non_stream: {e!r}", file=sys.stderr)
|
||||
|
||||
if HAS_TOOL_PARSERS and self.tool_parser_name and request.Tools:
|
||||
try:
|
||||
tools_raw = json.loads(request.Tools)
|
||||
tools = [SglTool.model_validate(t) for t in tools_raw] if SglTool else tools_raw
|
||||
fresh_tool_parser = FunctionCallParser(
|
||||
tools=tools, tool_call_parser=self.tool_parser_name,
|
||||
)
|
||||
normal, calls = fresh_tool_parser.parse_non_stream(final_content)
|
||||
if calls:
|
||||
final_content = normal
|
||||
for tc in calls:
|
||||
idx = int(getattr(tc, "tool_index", 0) or 0)
|
||||
final_tool_calls.append(backend_pb2.ToolCallDelta(
|
||||
index=idx,
|
||||
id=f"call_{uuid.uuid4().hex[:24]}",
|
||||
name=getattr(tc, "name", "") or "",
|
||||
arguments=getattr(tc, "parameters", "") or "",
|
||||
))
|
||||
except Exception as e:
|
||||
print(f"tool_parser.parse_non_stream: {e!r}", file=sys.stderr)
|
||||
|
||||
chat_delta = backend_pb2.ChatDelta(
|
||||
content=final_content if not streaming else "",
|
||||
reasoning_content=final_reasoning,
|
||||
tool_calls=final_tool_calls,
|
||||
)
|
||||
|
||||
if streaming:
|
||||
yield backend_pb2.Reply(
|
||||
message=b"",
|
||||
prompt_tokens=prompt_tokens,
|
||||
tokens=completion_tokens,
|
||||
chat_deltas=[chat_delta],
|
||||
)
|
||||
return
|
||||
|
||||
yield backend_pb2.Reply(
|
||||
message=bytes(final_content or "", "utf-8"),
|
||||
prompt_tokens=prompt_tokens,
|
||||
tokens=completion_tokens,
|
||||
chat_deltas=[chat_delta],
|
||||
)
|
||||
|
||||
|
||||
async def serve(address):
|
||||
server = grpc.aio.server(
|
||||
migration_thread_pool=futures.ThreadPoolExecutor(max_workers=MAX_WORKERS),
|
||||
options=[
|
||||
('grpc.max_message_length', 50 * 1024 * 1024),
|
||||
('grpc.max_send_message_length', 50 * 1024 * 1024),
|
||||
('grpc.max_receive_message_length', 50 * 1024 * 1024),
|
||||
],
|
||||
interceptors=get_auth_interceptors(aio=True),
|
||||
)
|
||||
backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
|
||||
server.add_insecure_port(address)
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
for sig in (signal.SIGINT, signal.SIGTERM):
|
||||
loop.add_signal_handler(sig, lambda: asyncio.ensure_future(server.stop(5)))
|
||||
|
||||
await server.start()
|
||||
print("Server started. Listening on: " + address, file=sys.stderr)
|
||||
await server.wait_for_termination()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Run the sglang gRPC server.")
|
||||
parser.add_argument(
|
||||
"--addr", default="localhost:50051", help="The address to bind the server to.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
asyncio.run(serve(args.addr))
|
||||
87
backend/python/sglang/install.sh
Executable file
87
backend/python/sglang/install.sh
Executable file
@@ -0,0 +1,87 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
EXTRA_PIP_INSTALL_FLAGS="--no-build-isolation"
|
||||
|
||||
# Avoid overcommitting the CPU during builds that compile native code.
|
||||
export NVCC_THREADS=2
|
||||
export MAX_JOBS=1
|
||||
|
||||
backend_dir=$(dirname $0)
|
||||
|
||||
if [ -d $backend_dir/common ]; then
|
||||
source $backend_dir/common/libbackend.sh
|
||||
else
|
||||
source $backend_dir/../common/libbackend.sh
|
||||
fi
|
||||
|
||||
if [ "x${BUILD_PROFILE}" == "xintel" ]; then
|
||||
EXTRA_PIP_INSTALL_FLAGS+=" --upgrade --index-strategy=unsafe-first-match"
|
||||
fi
|
||||
|
||||
if [ "x${BUILD_PROFILE}" == "xcpu" ]; then
|
||||
EXTRA_PIP_INSTALL_FLAGS+=" --index-strategy=unsafe-best-match"
|
||||
fi
|
||||
|
||||
# sglang's CPU path has no prebuilt wheel on PyPI — upstream publishes
|
||||
# a separate pyproject_cpu.toml that must be swapped in before `pip install`.
|
||||
# Reference: docker/xeon.Dockerfile in the sglang upstream repo.
|
||||
#
|
||||
# When BUILD_TYPE is empty (CPU profile) or FROM_SOURCE=true is forced,
|
||||
# install torch/transformers/etc from requirements-cpu.txt, then clone
|
||||
# sglang and install its python/ and sgl-kernel/ packages from source
|
||||
# using the CPU pyproject.
|
||||
if [ "x${BUILD_TYPE}" == "x" ] || [ "x${FROM_SOURCE:-}" == "xtrue" ]; then
|
||||
# sgl-kernel's CPU build links against libnuma and libtbb. Install
|
||||
# them here (Docker builder stage) before running the source build.
|
||||
# Harmless no-op on runs outside the docker build since installRequirements
|
||||
# below still needs them only if we reach the source build branch.
|
||||
if command -v apt-get >/dev/null 2>&1 && [ "$(id -u)" = "0" ]; then
|
||||
apt-get update
|
||||
DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
|
||||
libnuma-dev numactl libtbb-dev libgomp1 libomp-dev google-perftools \
|
||||
build-essential cmake ninja-build
|
||||
fi
|
||||
|
||||
installRequirements
|
||||
|
||||
# sgl-kernel's pyproject_cpu.toml uses scikit-build-core as its build
|
||||
# backend. With --no-build-isolation, that (and ninja/cmake) must be
|
||||
# present in the venv before we build from source.
|
||||
uv pip install --no-build-isolation "scikit-build-core>=0.10" ninja cmake
|
||||
|
||||
# sgl-kernel's CPU shm.cpp uses __m512 AVX-512 intrinsics unconditionally.
|
||||
# csrc/cpu/CMakeLists.txt hard-codes add_compile_options(-march=native),
|
||||
# which on runners without AVX-512 in /proc/cpuinfo fails with
|
||||
# "__m512 return without 'avx512f' enabled changes the ABI".
|
||||
# CXXFLAGS alone is insufficient because CMake's add_compile_options()
|
||||
# appends -march=native *after* CXXFLAGS, overriding it.
|
||||
# We therefore patch the CMakeLists.txt to replace -march=native with
|
||||
# -march=sapphirerapids so the flag is consistent throughout the build.
|
||||
# The resulting binary still requires an AVX-512 capable CPU at runtime,
|
||||
# same constraint sglang upstream documents in docker/xeon.Dockerfile.
|
||||
|
||||
_sgl_src=$(mktemp -d)
|
||||
trap 'rm -rf "${_sgl_src}"' EXIT
|
||||
git clone --depth 1 https://github.com/sgl-project/sglang "${_sgl_src}/sglang"
|
||||
|
||||
# Patch -march=native → -march=sapphirerapids in the CPU kernel CMakeLists
|
||||
sed -i 's/-march=native/-march=sapphirerapids/g' \
|
||||
"${_sgl_src}/sglang/sgl-kernel/csrc/cpu/CMakeLists.txt"
|
||||
|
||||
pushd "${_sgl_src}/sglang/sgl-kernel"
|
||||
if [ -f pyproject_cpu.toml ]; then
|
||||
cp pyproject_cpu.toml pyproject.toml
|
||||
fi
|
||||
uv pip install ${EXTRA_PIP_INSTALL_FLAGS:-} .
|
||||
popd
|
||||
|
||||
pushd "${_sgl_src}/sglang/python"
|
||||
if [ -f pyproject_cpu.toml ]; then
|
||||
cp pyproject_cpu.toml pyproject.toml
|
||||
fi
|
||||
uv pip install ${EXTRA_PIP_INSTALL_FLAGS:-} .
|
||||
popd
|
||||
else
|
||||
installRequirements
|
||||
fi
|
||||
63
backend/python/sglang/package.sh
Executable file
63
backend/python/sglang/package.sh
Executable file
@@ -0,0 +1,63 @@
|
||||
#!/bin/bash
|
||||
# Package runtime shared libraries for the sglang backend.
|
||||
#
|
||||
# Dockerfile.python's final stage is FROM scratch — every system library
|
||||
# the backend dlopens at runtime must be explicitly copied into
|
||||
# ${BACKEND}/lib, which libbackend.sh adds to LD_LIBRARY_PATH.
|
||||
#
|
||||
# sglang's CPU kernel links against libnuma and libtbb; torch's CPU
|
||||
# kernels use libgomp; tcmalloc + iomp5 are preloaded per sglang's
|
||||
# docker/xeon.Dockerfile recipe for best CPU throughput. Missing any of
|
||||
# these makes the engine crash on import.
|
||||
|
||||
set -e
|
||||
|
||||
CURDIR=$(dirname "$(realpath "$0")")
|
||||
LIB_DIR="${CURDIR}/lib"
|
||||
mkdir -p "${LIB_DIR}"
|
||||
|
||||
copy_with_symlinks() {
|
||||
local soname="$1"
|
||||
local hit=""
|
||||
for dir in \
|
||||
/usr/lib/x86_64-linux-gnu \
|
||||
/usr/lib/aarch64-linux-gnu \
|
||||
/lib/x86_64-linux-gnu \
|
||||
/lib/aarch64-linux-gnu \
|
||||
/usr/lib \
|
||||
/lib; do
|
||||
if [ -e "${dir}/${soname}" ]; then
|
||||
hit="${dir}/${soname}"
|
||||
break
|
||||
fi
|
||||
done
|
||||
if [ -z "${hit}" ]; then
|
||||
echo "warning: ${soname} not found in standard lib paths" >&2
|
||||
return 0
|
||||
fi
|
||||
local real
|
||||
real=$(readlink -f "${hit}")
|
||||
cp -v "${real}" "${LIB_DIR}/"
|
||||
local real_base
|
||||
real_base=$(basename "${real}")
|
||||
if [ "${real_base}" != "${soname}" ]; then
|
||||
ln -sf "${real_base}" "${LIB_DIR}/${soname}"
|
||||
fi
|
||||
}
|
||||
|
||||
copy_with_symlinks libnuma.so.1
|
||||
copy_with_symlinks libgomp.so.1
|
||||
copy_with_symlinks libtbb.so.12
|
||||
copy_with_symlinks libtbbmalloc.so.2
|
||||
copy_with_symlinks libtcmalloc.so.4
|
||||
|
||||
# intel-openmp ships libiomp5.so inside the venv under venv/lib/ — sglang's
|
||||
# CPU kernel was compiled against its __kmpc_* symbols, so it must be on
|
||||
# LD_LIBRARY_PATH at runtime. Copy it into the backend lib dir where
|
||||
# libbackend.sh will pick it up.
|
||||
if [ -f "${CURDIR}/venv/lib/libiomp5.so" ]; then
|
||||
cp -v "${CURDIR}/venv/lib/libiomp5.so" "${LIB_DIR}/"
|
||||
fi
|
||||
|
||||
echo "sglang packaging completed successfully"
|
||||
ls -liah "${LIB_DIR}/"
|
||||
2
backend/python/sglang/requirements-after.txt
Normal file
2
backend/python/sglang/requirements-after.txt
Normal file
@@ -0,0 +1,2 @@
|
||||
# sglang is installed per-acceleration in requirements-{profile}-after.txt
|
||||
# (cublas12, hipblas, intel, cpu)
|
||||
3
backend/python/sglang/requirements-cpu-after.txt
Normal file
3
backend/python/sglang/requirements-cpu-after.txt
Normal file
@@ -0,0 +1,3 @@
|
||||
# sglang has no prebuilt CPU wheel on PyPI. install.sh performs a
|
||||
# from-source build using the upstream pyproject_cpu.toml recipe from
|
||||
# docker/xeon.Dockerfile when BUILD_TYPE is empty (CPU profile).
|
||||
7
backend/python/sglang/requirements-cpu.txt
Normal file
7
backend/python/sglang/requirements-cpu.txt
Normal file
@@ -0,0 +1,7 @@
|
||||
--extra-index-url https://download.pytorch.org/whl/cpu
|
||||
accelerate
|
||||
torch==2.9.0
|
||||
torchvision
|
||||
torchaudio
|
||||
transformers
|
||||
intel-openmp; platform_machine == 'x86_64'
|
||||
3
backend/python/sglang/requirements-cublas12-after.txt
Normal file
3
backend/python/sglang/requirements-cublas12-after.txt
Normal file
@@ -0,0 +1,3 @@
|
||||
# Bump this pin deliberately — sglang releases weekly and API surfaces
|
||||
# (FunctionCallParser, ReasoningParser) move between releases.
|
||||
sglang[all]>=0.4.0
|
||||
5
backend/python/sglang/requirements-cublas12.txt
Normal file
5
backend/python/sglang/requirements-cublas12.txt
Normal file
@@ -0,0 +1,5 @@
|
||||
accelerate
|
||||
torch==2.7.1
|
||||
torchvision
|
||||
torchaudio==2.7.1
|
||||
transformers
|
||||
2
backend/python/sglang/requirements-hipblas-after.txt
Normal file
2
backend/python/sglang/requirements-hipblas-after.txt
Normal file
@@ -0,0 +1,2 @@
|
||||
# sglang's ROCm build is installed from source per docker/rocm.Dockerfile
|
||||
# upstream; install.sh handles the source build when BUILD_TYPE=hipblas.
|
||||
5
backend/python/sglang/requirements-hipblas.txt
Normal file
5
backend/python/sglang/requirements-hipblas.txt
Normal file
@@ -0,0 +1,5 @@
|
||||
--extra-index-url https://download.pytorch.org/whl/nightly/rocm7.0
|
||||
accelerate
|
||||
torch
|
||||
torchvision
|
||||
transformers
|
||||
6
backend/python/sglang/requirements-install.txt
Normal file
6
backend/python/sglang/requirements-install.txt
Normal file
@@ -0,0 +1,6 @@
|
||||
# sglang and sgl-kernel do not declare full PEP517 build deps; install the
|
||||
# basic build tooling into the venv before pulling the rest of the stack.
|
||||
packaging
|
||||
setuptools
|
||||
wheel
|
||||
setuptools-scm
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user