wehub/mudler--localai
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Security Scan / tests (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:12:26 +08:00
commit 1b8708893a
2967 changed files with 796875 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/python/README.md
+189
View File
@@ -0,0 +1,189 @@
# Python Backends for LocalAI
This directory contains Python-based AI backends for LocalAI, providing support for various AI models and hardware acceleration targets.
## Overview
The Python backends use a unified build system based on `libbackend.sh` that provides:
- **Automatic virtual environment management** with support for both `uv` and `pip`
- **Hardware-specific dependency installation** (CPU, CUDA, Intel, MLX, etc.)
- **Portable Python support** for standalone deployments
- **Consistent backend execution** across different environments
## Available Backends
### Core AI Models
- **transformers** - Hugging Face Transformers framework (PyTorch-based)
- **vllm** - High-performance LLM inference engine
- **mlx** - Apple Silicon optimized ML framework
### Audio & Speech
- **coqui** - Coqui TTS models
- **faster-whisper** - Fast Whisper speech recognition
- **kitten-tts** - Lightweight TTS
- **mlx-audio** - Apple Silicon audio processing
- **chatterbox** - TTS model
- **kokoro** - TTS models
### Computer Vision
- **diffusers** - Stable Diffusion and image generation
- **longcat-video** - CUDA video and speech-driven avatar generation with LongCat-Video
- **mlx-vlm** - Vision-language models for Apple Silicon
- **rfdetr** - Object detection models
### Specialized
- **rerankers** - Text reranking models
## Quick Start
### Prerequisites
- Python 3.10+ (default: 3.10.18)
- `uv` package manager (recommended) or `pip`
- Appropriate hardware drivers for your target (CUDA, Intel, etc.)
### Installation
Each backend can be installed individually:
```bash
# Navigate to a specific backend
cd backend/python/transformers
# Install dependencies
make transformers
# or
bash install.sh
# Run the backend
make run
# or
bash run.sh
```
### Using the Unified Build System
The `libbackend.sh` script provides consistent commands across all backends:
```bash
# Source the library in your backend script
source $(dirname $0)/../common/libbackend.sh
# Install requirements (automatically handles hardware detection)
installRequirements
# Start the backend server
startBackend $@
# Run tests
runUnittests
```
## Hardware Targets
The build system automatically detects and configures for different hardware:
- **CPU** - Standard CPU-only builds
- **CUDA** - NVIDIA GPU acceleration (supports CUDA 12/13)
- **Intel** - Intel XPU/GPU optimization
- **MLX** - Apple Silicon (M1/M2/M3) optimization
- **HIP** - AMD GPU acceleration
### Target-Specific Requirements
Backends can specify hardware-specific dependencies:
- `requirements.txt` - Base requirements
- `requirements-cpu.txt` - CPU-specific packages
- `requirements-cublas12.txt` - CUDA 12 packages
- `requirements-cublas13.txt` - CUDA 13 packages
- `requirements-intel.txt` - Intel-optimized packages
- `requirements-mps.txt` - Apple Silicon packages
## Configuration Options
### Environment Variables
- `PYTHON_VERSION` - Python version (default: 3.10)
- `PYTHON_PATCH` - Python patch version (default: 18)
- `BUILD_TYPE` - Force specific build target
- `USE_PIP` - Use pip instead of uv (default: false)
- `PORTABLE_PYTHON` - Enable portable Python builds
- `LIMIT_TARGETS` - Restrict backend to specific targets
### Example: CUDA 12 Only Backend
```bash
# In your backend script
LIMIT_TARGETS="cublas12"
source $(dirname $0)/../common/libbackend.sh
```
### Example: Intel-Optimized Backend
```bash
# In your backend script
LIMIT_TARGETS="intel"
source $(dirname $0)/../common/libbackend.sh
```
## Development
### Adding a New Backend
1. Create a new directory in `backend/python/`
2. Copy the template structure from `common/template/`
3. Implement your `backend.py` with the required gRPC interface
4. Add appropriate requirements files for your target hardware
5. Use `libbackend.sh` for consistent build and execution
### Testing
```bash
# Run backend tests
make test
# or
bash test.sh
```
### Building
```bash
# Install dependencies
make <backend-name>
# Clean build artifacts
make clean
```
## Architecture
Each backend follows a consistent structure:
```
backend-name/
├── backend.py # Main backend implementation
├── requirements.txt # Base dependencies
├── requirements-*.txt # Hardware-specific dependencies
├── install.sh # Installation script
├── run.sh # Execution script
├── test.sh # Test script
├── Makefile # Build targets
└── test.py # Unit tests
```
## Troubleshooting
### Common Issues
1. **Missing dependencies**: Ensure all requirements files are properly configured
2. **Hardware detection**: Check that `BUILD_TYPE` matches your system
3. **Python version**: Verify Python 3.10+ is available
4. **Virtual environment**: Use `ensureVenv` to create/activate environments
## Contributing
When adding new backends or modifying existing ones:
1. Follow the established directory structure
2. Use `libbackend.sh` for consistent behavior
3. Include appropriate requirements files for all target hardware
4. Add comprehensive tests
5. Update this README if adding new backend types
backend/python/ace-step/Makefile
+16
View 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
backend/python/ace-step/backend.py
+478
View File
@@ -0,0 +1,478 @@
#!/usr/bin/env python3
"""
LocalAI ACE-Step Backend
gRPC backend for ACE-Step 1.5 music generation. Aligns with upstream acestep API:
- LoadModel: initializes AceStepHandler (DiT) and LLMHandler, parses Options.
- SoundGeneration: uses create_sample (simple mode), format_sample (optional), then
generate_music from acestep.inference. Writes first output to request.dst.
- Fail hard: no fallback WAV on error; exceptions propagate to gRPC.
"""
from concurrent import futures
import argparse
import shutil
import signal
import sys
import os
import tempfile
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
from acestep.inference import (
GenerationParams,
GenerationConfig,
generate_music,
create_sample,
format_sample,
)
from acestep.handler import AceStepHandler
from acestep.llm_inference import LLMHandler
from acestep.model_downloader import ensure_lm_model
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
MAX_WORKERS = int(os.environ.get("PYTHON_GRPC_MAX_WORKERS", "1"))
# Model name -> HuggingFace/ModelScope repo (from upstream api_server.py)
MODEL_REPO_MAPPING = {
"acestep-v15-turbo": "ACE-Step/Ace-Step1.5",
"acestep-5Hz-lm-0.6B": "ACE-Step/Ace-Step1.5",
"acestep-5Hz-lm-1.7B": "ACE-Step/Ace-Step1.5",
"vae": "ACE-Step/Ace-Step1.5",
"Qwen3-Embedding-0.6B": "ACE-Step/Ace-Step1.5",
"acestep-v15-base": "ACE-Step/acestep-v15-base",
"acestep-v15-sft": "ACE-Step/acestep-v15-sft",
"acestep-v15-turbo-shift3": "ACE-Step/acestep-v15-turbo-shift3",
"acestep-5Hz-lm-4B": "ACE-Step/acestep-5Hz-lm-4B",
}
DEFAULT_REPO_ID = "ACE-Step/Ace-Step1.5"
def _is_float(s):
try:
float(s)
return True
except (ValueError, TypeError):
return False
def _is_int(s):
try:
int(s)
return True
except (ValueError, TypeError):
return False
def _parse_timesteps(s):
if s is None or (isinstance(s, str) and not s.strip()):
return None
if isinstance(s, (list, tuple)):
return [float(x) for x in s]
try:
return [float(x.strip()) for x in str(s).split(",") if x.strip()]
except (ValueError, TypeError):
return None
def _parse_options(opts_list):
"""Parse repeated 'key:value' options into a dict. Coerce numeric and bool."""
out = {}
for opt in opts_list or []:
if ":" not in opt:
continue
key, value = opt.split(":", 1)
key = key.strip()
value = value.strip()
if _is_int(value):
out[key] = int(value)
elif _is_float(value):
out[key] = float(value)
elif value.lower() in ("true", "false"):
out[key] = value.lower() == "true"
else:
out[key] = value
return out
def _generate_audio_sync(servicer, payload, dst_path):
"""
Run full ACE-Step pipeline using acestep.inference:
- If sample_mode/sample_query: create_sample() for caption/lyrics/metadata.
- If use_format and caption/lyrics: format_sample().
- Build GenerationParams and GenerationConfig, then generate_music().
Writes the first generated audio to dst_path. Raises on failure.
"""
opts = servicer.options
dit_handler = servicer.dit_handler
llm_handler = servicer.llm_handler
for key, value in opts.items():
if key not in payload:
payload[key] = value
def _opt(name, default):
return opts.get(name, default)
lm_temperature = _opt("temperature", 0.85)
lm_cfg_scale = _opt("lm_cfg_scale", _opt("cfg_scale", 2.0))
lm_top_k = opts.get("top_k")
lm_top_p = _opt("top_p", 0.9)
if lm_top_p is not None and lm_top_p >= 1.0:
lm_top_p = None
inference_steps = _opt("inference_steps", 8)
guidance_scale = _opt("guidance_scale", 7.0)
batch_size = max(1, int(_opt("batch_size", 1)))
use_simple = bool(payload.get("sample_query") or payload.get("text"))
sample_mode = use_simple and (payload.get("thinking") or payload.get("sample_mode"))
sample_query = (payload.get("sample_query") or payload.get("text") or "").strip()
use_format = bool(payload.get("use_format"))
caption = (payload.get("prompt") or payload.get("caption") or "").strip()
lyrics = (payload.get("lyrics") or "").strip()
vocal_language = (payload.get("vocal_language") or "en").strip()
instrumental = bool(payload.get("instrumental"))
bpm = payload.get("bpm")
key_scale = (payload.get("key_scale") or "").strip()
time_signature = (payload.get("time_signature") or "").strip()
audio_duration = payload.get("audio_duration")
if audio_duration is not None:
try:
audio_duration = float(audio_duration)
except (TypeError, ValueError):
audio_duration = None
if sample_mode and llm_handler and getattr(llm_handler, "llm_initialized", False):
parsed_language = None
if sample_query:
for hint in ("english", "en", "chinese", "zh", "japanese", "ja"):
if hint in sample_query.lower():
parsed_language = "en" if hint == "english" or hint == "en" else hint
break
vocal_lang = vocal_language if vocal_language and vocal_language != "unknown" else parsed_language
sample_result = create_sample(
llm_handler=llm_handler,
query=sample_query or "NO USER INPUT",
instrumental=instrumental,
vocal_language=vocal_lang,
temperature=lm_temperature,
top_k=lm_top_k,
top_p=lm_top_p,
use_constrained_decoding=True,
)
if not sample_result.success:
raise RuntimeError(f"create_sample failed: {sample_result.error or sample_result.status_message}")
caption = sample_result.caption or caption
lyrics = sample_result.lyrics or lyrics
bpm = sample_result.bpm
key_scale = sample_result.keyscale or key_scale
time_signature = sample_result.timesignature or time_signature
if sample_result.duration is not None:
audio_duration = sample_result.duration
if getattr(sample_result, "language", None):
vocal_language = sample_result.language
if use_format and (caption or lyrics) and llm_handler and getattr(llm_handler, "llm_initialized", False):
user_metadata = {}
if bpm is not None:
user_metadata["bpm"] = bpm
if audio_duration is not None and float(audio_duration) > 0:
user_metadata["duration"] = int(audio_duration)
if key_scale:
user_metadata["keyscale"] = key_scale
if time_signature:
user_metadata["timesignature"] = time_signature
if vocal_language and vocal_language != "unknown":
user_metadata["language"] = vocal_language
format_result = format_sample(
llm_handler=llm_handler,
caption=caption,
lyrics=lyrics,
user_metadata=user_metadata if user_metadata else None,
temperature=lm_temperature,
top_k=lm_top_k,
top_p=lm_top_p,
use_constrained_decoding=True,
)
if format_result.success:
caption = format_result.caption or caption
lyrics = format_result.lyrics or lyrics
if format_result.duration is not None:
audio_duration = format_result.duration
if format_result.bpm is not None:
bpm = format_result.bpm
if format_result.keyscale:
key_scale = format_result.keyscale
if format_result.timesignature:
time_signature = format_result.timesignature
if getattr(format_result, "language", None):
vocal_language = format_result.language
thinking = bool(payload.get("thinking"))
use_cot_metas = not sample_mode
params = GenerationParams(
task_type=payload.get("task_type", "text2music"),
instruction=payload.get("instruction", "Fill the audio semantic mask based on the given conditions:"),
reference_audio=payload.get("reference_audio_path"),
src_audio=payload.get("src_audio_path"),
audio_codes=payload.get("audio_code_string", ""),
caption=caption,
lyrics=lyrics,
instrumental=instrumental or (not lyrics or str(lyrics).strip().lower() in ("[inst]", "[instrumental]")),
vocal_language=vocal_language or "unknown",
bpm=bpm,
keyscale=key_scale,
timesignature=time_signature,
duration=float(audio_duration) if audio_duration and float(audio_duration) > 0 else -1.0,
inference_steps=inference_steps,
seed=int(payload.get("seed", -1)),
guidance_scale=guidance_scale,
use_adg=bool(payload.get("use_adg")),
cfg_interval_start=float(payload.get("cfg_interval_start", 0.0)),
cfg_interval_end=float(payload.get("cfg_interval_end", 1.0)),
shift=float(payload.get("shift", 1.0)),
infer_method=(payload.get("infer_method") or "ode").strip(),
timesteps=_parse_timesteps(payload.get("timesteps")),
repainting_start=float(payload.get("repainting_start", 0.0)),
repainting_end=float(payload.get("repainting_end", -1)) if payload.get("repainting_end") is not None else -1,
audio_cover_strength=float(payload.get("audio_cover_strength", 1.0)),
thinking=thinking,
lm_temperature=lm_temperature,
lm_cfg_scale=lm_cfg_scale,
lm_top_k=lm_top_k or 0,
lm_top_p=lm_top_p if lm_top_p is not None and lm_top_p < 1.0 else 0.9,
lm_negative_prompt=payload.get("lm_negative_prompt", "NO USER INPUT"),
use_cot_metas=use_cot_metas,
use_cot_caption=bool(payload.get("use_cot_caption", True)),
use_cot_language=bool(payload.get("use_cot_language", True)),
use_constrained_decoding=True,
)
config = GenerationConfig(
batch_size=batch_size,
allow_lm_batch=bool(payload.get("allow_lm_batch", False)),
use_random_seed=bool(payload.get("use_random_seed", True)),
seeds=payload.get("seeds"),
lm_batch_chunk_size=max(1, int(payload.get("lm_batch_chunk_size", 8))),
constrained_decoding_debug=bool(payload.get("constrained_decoding_debug")),
audio_format=(payload.get("audio_format") or "flac").strip() or "flac",
)
save_dir = tempfile.mkdtemp(prefix="ace_step_")
try:
result = generate_music(
dit_handler=dit_handler,
llm_handler=llm_handler if (llm_handler and getattr(llm_handler, "llm_initialized", False)) else None,
params=params,
config=config,
save_dir=save_dir,
progress=None,
)
if not result.success:
raise RuntimeError(result.error or result.status_message or "generate_music failed")
audios = result.audios or []
if not audios:
raise RuntimeError("generate_music returned no audio")
first_path = audios[0].get("path") or ""
if not first_path or not os.path.isfile(first_path):
raise RuntimeError("first generated audio path missing or not a file")
shutil.copy2(first_path, dst_path)
finally:
try:
shutil.rmtree(save_dir, ignore_errors=True)
except Exception:
pass
class BackendServicer(backend_pb2_grpc.BackendServicer):
def __init__(self):
self.model_path = None
self.model_dir = None
self.checkpoint_dir = None
self.project_root = None
self.options = {}
self.dit_handler = None
self.llm_handler = None
def Health(self, request, context):
return backend_pb2.Reply(message=b"OK")
def LoadModel(self, request, context):
try:
self.options = _parse_options(list(getattr(request, "Options", []) or []))
model_path = getattr(request, "ModelPath", None) or ""
model_name = (request.Model or "").strip()
model_file = (getattr(request, "ModelFile", None) or "").strip()
# Model dir: where we store checkpoints (always under LocalAI models path, never backend dir)
if model_path and model_name:
model_dir = os.path.join(model_path, model_name)
elif model_file:
model_dir = model_file
else:
model_dir = os.path.abspath(model_name or ".")
self.model_dir = model_dir
self.checkpoint_dir = os.path.join(model_dir, "checkpoints")
self.project_root = model_dir
self.model_path = os.path.join(self.checkpoint_dir, model_name or os.path.basename(model_dir.rstrip("/\\")))
config_path = model_name or os.path.basename(model_dir.rstrip("/\\"))
os.makedirs(self.checkpoint_dir, exist_ok=True)
self.dit_handler = AceStepHandler()
# Patch handler so it uses our model dir instead of site-packages/checkpoints
self.dit_handler._get_project_root = lambda: self.project_root
device = self.options.get("device", "auto")
use_flash = self.options.get("use_flash_attention", True)
if isinstance(use_flash, str):
use_flash = str(use_flash).lower() in ("1", "true", "yes")
offload = self.options.get("offload_to_cpu", False)
if isinstance(offload, str):
offload = str(offload).lower() in ("1", "true", "yes")
status_msg, ok = self.dit_handler.initialize_service(
project_root=self.project_root,
config_path=config_path,
device=device,
use_flash_attention=use_flash,
compile_model=False,
offload_to_cpu=offload,
offload_dit_to_cpu=bool(self.options.get("offload_dit_to_cpu", False)),
)
if not ok:
return backend_pb2.Result(success=False, message=f"DiT init failed: {status_msg}")
self.llm_handler = None
if self.options.get("init_lm", True):
lm_model = self.options.get("lm_model_path", "acestep-5Hz-lm-0.6B")
# Ensure LM model is downloaded before initializing
try:
from pathlib import Path
lm_success, lm_msg = ensure_lm_model(
model_name=lm_model,
checkpoints_dir=Path(self.checkpoint_dir),
prefer_source=None, # Auto-detect HuggingFace vs ModelScope
)
if not lm_success:
print(f"[ace-step] Warning: LM model download failed: {lm_msg}", file=sys.stderr)
# Continue anyway - LLM initialization will fail gracefully
else:
print(f"[ace-step] LM model ready: {lm_msg}", file=sys.stderr)
except Exception as e:
print(f"[ace-step] Warning: LM model download check failed: {e}", file=sys.stderr)
# Continue anyway - LLM initialization will fail gracefully
self.llm_handler = LLMHandler()
lm_backend = (self.options.get("lm_backend") or "vllm").strip().lower()
if lm_backend not in ("vllm", "pt"):
lm_backend = "vllm"
lm_status, lm_ok = self.llm_handler.initialize(
checkpoint_dir=self.checkpoint_dir,
lm_model_path=lm_model,
backend=lm_backend,
device=device,
offload_to_cpu=offload,
dtype=getattr(self.dit_handler, "dtype", None),
)
if not lm_ok:
self.llm_handler = None
print(f"[ace-step] LM init failed (optional): {lm_status}", file=sys.stderr)
print(f"[ace-step] LoadModel: model={self.model_path}, options={list(self.options.keys())}", file=sys.stderr)
return backend_pb2.Result(success=True, message="Model loaded successfully")
except Exception as err:
return backend_pb2.Result(success=False, message=f"LoadModel error: {err}")
def SoundGeneration(self, request, context):
if not request.dst:
return backend_pb2.Result(success=False, message="request.dst is required")
use_simple = bool(request.text)
if use_simple:
payload = {
"sample_query": request.text or "",
"sample_mode": True,
"thinking": True,
"vocal_language": request.language or request.GetLanguage() or "en",
"instrumental": request.instrumental if request.HasField("instrumental") else False,
}
else:
caption = request.caption or request.GetCaption() or request.text
payload = {
"prompt": caption,
"lyrics": request.lyrics or request.lyrics or "",
"thinking": request.think if request.HasField("think") else False,
"vocal_language": request.language or request.GetLanguage() or "en",
}
if request.HasField("bpm"):
payload["bpm"] = request.bpm
if request.HasField("keyscale") and request.keyscale:
payload["key_scale"] = request.keyscale
if request.HasField("timesignature") and request.timesignature:
payload["time_signature"] = request.timesignature
if request.HasField("duration") and request.duration:
payload["audio_duration"] = int(request.duration) if request.duration else None
if request.src:
payload["src_audio_path"] = request.src
_generate_audio_sync(self, payload, request.dst)
return backend_pb2.Result(success=True, message="Sound generated successfully")
def TTS(self, request, context):
if not request.dst:
return backend_pb2.Result(success=False, message="request.dst is required")
payload = {
"sample_query": request.text,
"sample_mode": True,
"thinking": False,
"vocal_language": (request.language if request.language else "") or "en",
"instrumental": False,
}
_generate_audio_sync(self, payload, request.dst)
return backend_pb2.Result(success=True, message="TTS (music fallback) generated successfully")
def serve(address):
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(f"[ace-step] Server listening on {address}", file=sys.stderr)
def shutdown(sig, frame):
server.stop(0)
sys.exit(0)
signal.signal(signal.SIGINT, shutdown)
signal.signal(signal.SIGTERM, shutdown)
try:
while True:
import time
time.sleep(_ONE_DAY_IN_SECONDS)
except KeyboardInterrupt:
server.stop(0)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--addr", default="localhost:50051", help="Listen address")
args = parser.parse_args()
serve(args.addr)
backend/python/ace-step/install.sh
+26
View File
@@ -0,0 +1,26 @@
#!/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
PYTHON_VERSION="3.11"
PYTHON_PATCH="14"
PY_STANDALONE_TAG="20260203"
installRequirements
if [ ! -d ACE-Step-1.5 ]; then
git clone https://github.com/ace-step/ACE-Step-1.5
cd ACE-Step-1.5/
if [ "x${USE_PIP}" == "xtrue" ]; then
pip install ${EXTRA_PIP_INSTALL_FLAGS:-} --no-deps .
else
uv pip install ${EXTRA_PIP_INSTALL_FLAGS:-} --no-deps .
fi
fi
backend/python/ace-step/requirements-cpu.txt
+22
View File
@@ -0,0 +1,22 @@
--extra-index-url https://download.pytorch.org/whl/cpu
torch
torchaudio
torchvision
# Core dependencies
transformers>=4.51.0,<4.58.0
diffusers
gradio
matplotlib>=3.7.5
scipy>=1.10.1
soundfile>=0.13.1
loguru>=0.7.3
einops>=0.8.1
accelerate>=1.12.0
fastapi>=0.110.0
uvicorn[standard]>=0.27.0
numba>=0.63.1
vector-quantize-pytorch>=1.27.15
torchcodec>=0.9.1
torchao
modelscope
backend/python/ace-step/requirements-cublas12.txt
+22
View File
@@ -0,0 +1,22 @@
--extra-index-url https://download.pytorch.org/whl/cu128
torch
torchaudio
torchvision
# Core dependencies
transformers>=4.51.0,<4.58.0
diffusers
gradio>=6.5.1
matplotlib>=3.7.5
scipy>=1.10.1
soundfile>=0.13.1
loguru>=0.7.3
einops>=0.8.1
accelerate>=1.12.0
fastapi>=0.110.0
uvicorn[standard]>=0.27.0
numba>=0.63.1
vector-quantize-pytorch>=1.27.15
torchcodec>=0.9.1
torchao
modelscope
backend/python/ace-step/requirements-cublas13.txt
+22
View File
@@ -0,0 +1,22 @@
--extra-index-url https://download.pytorch.org/whl/cu130
torch
torchaudio
torchvision
# Core dependencies
transformers>=4.51.0,<4.58.0
diffusers
gradio>=6.5.1
matplotlib>=3.7.5
scipy>=1.10.1
soundfile>=0.13.1
loguru>=0.7.3
einops>=0.8.1
accelerate>=1.12.0
fastapi>=0.110.0
uvicorn[standard]>=0.27.0
numba>=0.63.1
vector-quantize-pytorch>=1.27.15
torchcodec>=0.9.1
torchao
modelscope
backend/python/ace-step/requirements-hipblas.txt
+22
View File
@@ -0,0 +1,22 @@
--extra-index-url https://download.pytorch.org/whl/rocm7.0
torch==2.10.0+rocm7.0
torchaudio
torchvision
# Core dependencies
transformers>=4.51.0,<4.58.0
diffusers
gradio>=6.5.1
matplotlib>=3.7.5
scipy>=1.10.1
soundfile>=0.13.1
loguru>=0.7.3
einops>=0.8.1
accelerate>=1.12.0
fastapi>=0.110.0
uvicorn[standard]>=0.27.0
numba>=0.63.1
vector-quantize-pytorch>=1.27.15
torchcodec>=0.9.1
torchao
modelscope
backend/python/ace-step/requirements-intel.txt
+26
View File
@@ -0,0 +1,26 @@
--extra-index-url https://download.pytorch.org/whl/xpu
torch
torchaudio
torchvision
# Core dependencies
transformers>=4.51.0,<4.58.0
diffusers
gradio
matplotlib>=3.7.5
scipy>=1.10.1
soundfile>=0.13.1
loguru>=0.7.3
einops>=0.8.1
accelerate>=1.12.0
fastapi>=0.110.0
uvicorn[standard]>=0.27.0
numba>=0.63.1
vector-quantize-pytorch>=1.27.15
torchcodec>=0.9.1
torchao
modelscope
# LoRA Training dependencies (optional)
peft>=0.7.0
lightning>=2.0.0
backend/python/ace-step/requirements-l4t13.txt
+21
View File
@@ -0,0 +1,21 @@
--extra-index-url https://download.pytorch.org/whl/cu130
torch
torchaudio
torchvision
# Core dependencies
transformers>=4.51.0,<4.58.0
diffusers
gradio>=6.5.1
matplotlib>=3.7.5
scipy>=1.10.1
soundfile>=0.13.1
loguru>=0.7.3
einops>=0.8.1
accelerate>=1.12.0
fastapi>=0.110.0
uvicorn[standard]>=0.27.0
numba>=0.63.1
vector-quantize-pytorch>=1.27.15
torchcodec>=0.9.1
torchao
modelscope
backend/python/ace-step/requirements-mps.txt
+25
View File
@@ -0,0 +1,25 @@
torch
torchaudio
torchvision
# Core dependencies
transformers>=4.51.0,<4.58.0
diffusers
gradio
matplotlib>=3.7.5
scipy>=1.10.1
soundfile>=0.13.1
loguru>=0.7.3
einops>=0.8.1
accelerate>=1.12.0
fastapi>=0.110.0
uvicorn[standard]>=0.27.0
numba>=0.63.1
vector-quantize-pytorch>=1.27.15
torchcodec>=0.9.1
torchao
modelscope
# LoRA Training dependencies (optional)
peft>=0.7.0
lightning>=2.0.0
backend/python/ace-step/requirements.txt
+4
View File
@@ -0,0 +1,4 @@
setuptools
grpcio==1.76.0
protobuf
certifi
backend/python/ace-step/run.sh
+9
View 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 $@
backend/python/ace-step/test.py
+53
View File
@@ -0,0 +1,53 @@
"""
Tests for the ACE-Step gRPC backend.
"""
import os
import tempfile
import unittest
import backend_pb2
import backend_pb2_grpc
import grpc
class TestACEStepBackend(unittest.TestCase):
"""Test Health, LoadModel, and SoundGeneration (minimal; no real model required)."""
@classmethod
def setUpClass(cls):
port = os.environ.get("BACKEND_PORT", "50051")
cls.channel = grpc.insecure_channel(f"localhost:{port}")
cls.stub = backend_pb2_grpc.BackendStub(cls.channel)
@classmethod
def tearDownClass(cls):
cls.channel.close()
def test_health(self):
response = self.stub.Health(backend_pb2.HealthMessage())
self.assertEqual(response.message, b"OK")
def test_load_model(self):
response = self.stub.LoadModel(backend_pb2.ModelOptions(Model="ace-step-test"))
self.assertTrue(response.success, response.message)
def test_sound_generation_minimal(self):
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
dst = f.name
try:
req = backend_pb2.SoundGenerationRequest(
text="upbeat pop song",
model="ace-step-test",
dst=dst,
)
response = self.stub.SoundGeneration(req)
self.assertTrue(response.success, response.message)
self.assertTrue(os.path.exists(dst), f"Output file not created: {dst}")
self.assertGreater(os.path.getsize(dst), 0)
finally:
if os.path.exists(dst):
os.unlink(dst)
if __name__ == "__main__":
unittest.main()
backend/python/ace-step/test.sh
+19
View File
@@ -0,0 +1,19 @@
#!/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
# Start backend in background (use env to avoid port conflict in parallel tests)
export PYTHONUNBUFFERED=1
BACKEND_PORT=${BACKEND_PORT:-50051}
python backend.py --addr "localhost:${BACKEND_PORT}" &
BACKEND_PID=$!
trap "kill $BACKEND_PID 2>/dev/null || true" EXIT
sleep 3
export BACKEND_PORT
runUnittests
backend/python/chatterbox/Makefile
+23
View File
@@ -0,0 +1,23 @@
.PHONY: chatterbox
chatterbox:
bash install.sh
.PHONY: run
run: chatterbox
@echo "Running coqui..."
bash run.sh
@echo "coqui run."
.PHONY: test
test: chatterbox
@echo "Testing coqui..."
bash test.sh
@echo "coqui tested."
.PHONY: protogen-clean
protogen-clean:
$(RM) backend_pb2_grpc.py backend_pb2.py
.PHONY: clean
clean: protogen-clean
rm -rf venv __pycache__
backend/python/chatterbox/backend.py
+285
View File
@@ -0,0 +1,285 @@
#!/usr/bin/env python3
"""
This is an extra gRPC server of LocalAI for Chatterbox TTS
"""
from concurrent import futures
import time
import argparse
import signal
import sys
import os
import backend_pb2
import backend_pb2_grpc
import torch
import torchaudio as ta
from chatterbox.tts import ChatterboxTTS
from chatterbox.mtl_tts import ChatterboxMultilingualTTS
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
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
def coerce_param_value(value):
"""Coerce a TTSRequest.params value (string on the wire) to the type the
Chatterbox generate() kwargs expect (float/int/bool), matching how static
YAML options are coerced at load time. Non-string values pass through."""
if not isinstance(value, str):
return value
if is_float(value):
return float(value)
if is_int(value):
return int(value)
if value.lower() in ["true", "false"]:
return value.lower() == "true"
return value
def split_text_at_word_boundary(text, max_length=250):
"""
Split text at word boundaries without truncating words.
Returns a list of text chunks.
"""
if not text or len(text) <= max_length:
return [text]
chunks = []
words = text.split()
current_chunk = ""
for word in words:
# Check if adding this word would exceed the limit
if len(current_chunk) + len(word) + 1 <= max_length:
if current_chunk:
current_chunk += " " + word
else:
current_chunk = word
else:
# If current chunk is not empty, add it to chunks
if current_chunk:
chunks.append(current_chunk)
current_chunk = word
else:
# If a single word is longer than max_length, we have to include it anyway
chunks.append(word)
current_chunk = ""
# Add the last chunk if it's not empty
if current_chunk:
chunks.append(current_chunk)
return chunks
def merge_audio_files(audio_files, output_path, sample_rate):
"""
Merge multiple audio files into a single audio file.
"""
if not audio_files:
return
if len(audio_files) == 1:
# If only one file, just copy it
import shutil
shutil.copy2(audio_files[0], output_path)
return
# Load all audio files
waveforms = []
for audio_file in audio_files:
waveform, sr = ta.load(audio_file)
if sr != sample_rate:
# Resample if necessary
resampler = ta.transforms.Resample(sr, sample_rate)
waveform = resampler(waveform)
waveforms.append(waveform)
# Concatenate all waveforms
merged_waveform = torch.cat(waveforms, dim=1)
# Save the merged audio
ta.save(output_path, merged_waveform, sample_rate)
# Clean up temporary files
for audio_file in audio_files:
if os.path.exists(audio_file):
os.remove(audio_file)
_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'))
# Implement the BackendServicer class with the service methods
class BackendServicer(backend_pb2_grpc.BackendServicer):
"""
BackendServicer is the class that implements the gRPC service
"""
def Health(self, request, context):
return backend_pb2.Reply(message=bytes("OK", 'utf-8'))
def LoadModel(self, request, context):
# Get device
# device = "cuda" if request.CUDA else "cpu"
if torch.cuda.is_available():
print("CUDA is available", file=sys.stderr)
device = "cuda"
else:
print("CUDA is not available", file=sys.stderr)
device = "cpu"
mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
if mps_available:
device = "mps"
if not torch.cuda.is_available() and request.CUDA:
return backend_pb2.Result(success=False, message="CUDA is not available")
options = request.Options
# empty dict
self.options = {}
# The options are a list of strings in this form optname:optvalue
# We are storing all the options in a dict so we can use it later when
# generating the images
for opt in options:
if ":" not in opt:
continue
key, value = opt.split(":")
# if value is a number, convert it to the appropriate type
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
self.AudioPath = None
if os.path.isabs(request.AudioPath):
self.AudioPath = request.AudioPath
elif request.AudioPath and request.ModelFile != "" and not os.path.isabs(request.AudioPath):
# get base path of modelFile
modelFileBase = os.path.dirname(request.ModelFile)
# modify LoraAdapter to be relative to modelFileBase
self.AudioPath = os.path.join(modelFileBase, request.AudioPath)
try:
print("Preparing models, please wait", file=sys.stderr)
if "multilingual" in self.options:
# remove key from options
del self.options["multilingual"]
self.model = ChatterboxMultilingualTTS.from_pretrained(device=device)
else:
self.model = ChatterboxTTS.from_pretrained(device=device)
except Exception as err:
return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")
# Implement your logic here for the LoadModel service
# Replace this with your desired response
return backend_pb2.Result(message="Model loaded successfully", success=True)
def TTS(self, request, context):
try:
kwargs = {}
if "language" in self.options:
kwargs["language_id"] = self.options["language"]
if self.AudioPath is not None:
kwargs["audio_prompt_path"] = self.AudioPath
# add options to kwargs
kwargs.update(self.options)
# Merge per-request params (TTSRequest.params), overriding the static
# YAML options. This exposes Chatterbox generation knobs (e.g.
# exaggeration, cfg_weight, temperature) per request. Values arrive as
# strings on the wire and are coerced to float/int/bool.
if hasattr(request, "params") and request.params:
for key, value in request.params.items():
kwargs[key] = coerce_param_value(value)
# Check if text exceeds 250 characters
# (chatterbox does not support long text)
# https://github.com/resemble-ai/chatterbox/issues/60
# https://github.com/resemble-ai/chatterbox/issues/110
if len(request.text) > 250:
# Split text at word boundaries
text_chunks = split_text_at_word_boundary(request.text, max_length=250)
print(f"Splitting text into chunks of 250 characters: {len(text_chunks)}", file=sys.stderr)
# Generate audio for each chunk
temp_audio_files = []
for i, chunk in enumerate(text_chunks):
# Generate audio for this chunk
wav = self.model.generate(chunk, **kwargs)
# Create temporary file for this chunk
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.wav')
temp_file.close()
ta.save(temp_file.name, wav, self.model.sr)
temp_audio_files.append(temp_file.name)
# Merge all audio files
merge_audio_files(temp_audio_files, request.dst, self.model.sr)
else:
# Generate audio using ChatterboxTTS for short text
wav = self.model.generate(request.text, **kwargs)
# Save the generated audio
ta.save(request.dst, wav, self.model.sr)
except Exception as err:
return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")
return backend_pb2.Result(success=True)
def serve(address):
server = grpc.server(futures.ThreadPoolExecutor(max_workers=MAX_WORKERS),
options=[
('grpc.max_message_length', 50 * 1024 * 1024), # 50MB
('grpc.max_send_message_length', 50 * 1024 * 1024), # 50MB
('grpc.max_receive_message_length', 50 * 1024 * 1024), # 50MB
],
interceptors=get_auth_interceptors(),
)
backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
server.add_insecure_port(address)
server.start()
print("Server started. Listening on: " + address, file=sys.stderr)
# Define the signal handler function
def signal_handler(sig, frame):
print("Received termination signal. Shutting down...")
server.stop(0)
sys.exit(0)
# Set the signal handlers for SIGINT and SIGTERM
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
try:
while True:
time.sleep(_ONE_DAY_IN_SECONDS)
except KeyboardInterrupt:
server.stop(0)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Run the gRPC server.")
parser.add_argument(
"--addr", default="localhost:50051", help="The address to bind the server to."
)
args = parser.parse_args()
serve(args.addr)
backend/python/chatterbox/install.sh
Executable
+37
View File
@@ -0,0 +1,37 @@
#!/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
# This is here because the Intel pip index is broken and returns 200 status codes for every package name, it just doesn't return any package links.
# This makes uv think that the package exists in the Intel pip index, and by default it stops looking at other pip indexes once it finds a match.
# We need uv to continue falling through to the pypi default index to find optimum[openvino] in the pypi index
# the --upgrade actually allows us to *downgrade* torch to the version provided in the Intel pip index
if [ "x${BUILD_PROFILE}" == "xintel" ]; then
EXTRA_PIP_INSTALL_FLAGS+=" --upgrade --index-strategy=unsafe-first-match"
fi
EXTRA_PIP_INSTALL_FLAGS+=" --no-build-isolation"
if [ "x${BUILD_PROFILE}" == "xl4t12" ]; then
USE_PIP=true
fi
installRequirements
# chatterbox-tts upstream pulls `russian-text-stresser` (unpinned git URL) which
# transitively pins spacy==3.6.* and other ancient packages. That cascade forces
# pip to backtrack through Jinja2/MarkupSafe/omegaconf/ruamel.yaml into Python-2-era
# sdists that no longer build. We install chatterbox-tts itself with --no-deps and
# list its real runtime deps in requirements-*.txt instead.
echo "Installing chatterbox-tts with --no-deps"
if [ "x${USE_PIP}" == "xtrue" ]; then
pip install ${EXTRA_PIP_INSTALL_FLAGS:-} --no-deps "chatterbox-tts@git+https://git@github.com/mudler/chatterbox.git@faster"
else
uv pip install ${EXTRA_PIP_INSTALL_FLAGS:-} --no-deps "chatterbox-tts@git+https://git@github.com/mudler/chatterbox.git@faster"
fi
backend/python/chatterbox/requirements-cpu.txt
+19
View File
@@ -0,0 +1,19 @@
--extra-index-url https://download.pytorch.org/whl/cpu
accelerate
torch
torchaudio
numpy>=1.24.0,<1.26.0
transformers
# chatterbox-tts itself is installed with --no-deps in install.sh.
# These are its real runtime deps, mirroring upstream's pyproject.toml
# minus russian-text-stresser (whose ancient pins break the resolver).
omegaconf==2.3.0
resampy==0.4.3
librosa
s3tokenizer
diffusers
resemble-perth==1.0.1
conformer
safetensors
spacy-pkuseg
pykakasi==2.3.0
backend/python/chatterbox/requirements-cublas12.txt
+18
View File
@@ -0,0 +1,18 @@
torch
torchaudio
transformers
numpy>=1.24.0,<1.26.0
# chatterbox-tts itself is installed with --no-deps in install.sh.
# These are its real runtime deps, mirroring upstream's pyproject.toml
# minus russian-text-stresser (whose ancient pins break the resolver).
omegaconf==2.3.0
resampy==0.4.3
librosa
s3tokenizer
diffusers
resemble-perth==1.0.1
conformer
safetensors
spacy-pkuseg
pykakasi==2.3.0
accelerate
backend/python/chatterbox/requirements-cublas13.txt
+19
View File
@@ -0,0 +1,19 @@
--extra-index-url https://download.pytorch.org/whl/cu130
torch
torchaudio
transformers
numpy>=1.24.0,<1.26.0
# chatterbox-tts itself is installed with --no-deps in install.sh.
# These are its real runtime deps, mirroring upstream's pyproject.toml
# minus russian-text-stresser (whose ancient pins break the resolver).
omegaconf==2.3.0
resampy==0.4.3
librosa
s3tokenizer
diffusers
resemble-perth==1.0.1
conformer
safetensors
spacy-pkuseg
pykakasi==2.3.0
accelerate
backend/python/chatterbox/requirements-hipblas.txt
+19
View File
@@ -0,0 +1,19 @@
--extra-index-url https://download.pytorch.org/whl/rocm7.0
torch==2.10.0+rocm7.0
torchaudio==2.10.0+rocm7.0
transformers
numpy>=1.24.0,<1.26.0
# chatterbox-tts itself is installed with --no-deps in install.sh.
# These are its real runtime deps, mirroring upstream's pyproject.toml
# minus russian-text-stresser (whose ancient pins break the resolver).
omegaconf==2.3.0
resampy==0.4.3
librosa
s3tokenizer
diffusers
resemble-perth==1.0.1
conformer
safetensors
spacy-pkuseg
pykakasi==2.3.0
accelerate
backend/python/chatterbox/requirements-install.txt
+5
View File
@@ -0,0 +1,5 @@
# Build dependencies needed for packages installed from source (e.g., git dependencies)
# When using --no-build-isolation, these must be installed in the venv first
wheel
setuptools
packaging
backend/python/chatterbox/requirements-intel.txt
+22
View File
@@ -0,0 +1,22 @@
--extra-index-url https://download.pytorch.org/whl/xpu
torch
torchaudio
transformers
numpy>=1.24.0,<1.26.0
# chatterbox-tts itself is installed with --no-deps in install.sh.
# These are its real runtime deps, mirroring upstream's pyproject.toml
# minus russian-text-stresser (whose ancient pins break the resolver).
omegaconf==2.3.0
resampy==0.4.3
librosa
s3tokenizer
diffusers
resemble-perth==1.0.1
conformer
safetensors
spacy-pkuseg
pykakasi==2.3.0
accelerate
oneccl_bind_pt==2.3.100+xpu
optimum[openvino]
setuptools
backend/python/chatterbox/requirements-l4t12.txt
+19
View File
@@ -0,0 +1,19 @@
--extra-index-url https://pypi.jetson-ai-lab.io/jp6/cu126/
torch
torchaudio
transformers
numpy>=1.24.0,<1.26.0
# chatterbox-tts itself is installed with --no-deps in install.sh.
# These are its real runtime deps, mirroring upstream's pyproject.toml
# minus russian-text-stresser (whose ancient pins break the resolver).
omegaconf==2.3.0
resampy==0.4.3
librosa
s3tokenizer
diffusers
resemble-perth==1.0.1
conformer
safetensors
spacy-pkuseg
pykakasi==2.3.0
accelerate
backend/python/chatterbox/requirements-l4t13.txt
+19
View File
@@ -0,0 +1,19 @@
--extra-index-url https://download.pytorch.org/whl/cu130
torch
torchaudio
transformers
numpy>=1.24.0,<1.26.0
# chatterbox-tts itself is installed with --no-deps in install.sh.
# These are its real runtime deps, mirroring upstream's pyproject.toml
# minus russian-text-stresser (whose ancient pins break the resolver).
omegaconf==2.3.0
resampy==0.4.3
librosa
s3tokenizer
diffusers
resemble-perth==1.0.1
conformer
safetensors
spacy-pkuseg
pykakasi==2.3.0
accelerate
backend/python/chatterbox/requirements-mps.txt
+18
View File
@@ -0,0 +1,18 @@
torch
torchaudio
accelerate
numpy>=1.24.0,<1.26.0
transformers
# chatterbox-tts itself is installed with --no-deps in install.sh.
# These are its real runtime deps, mirroring upstream's pyproject.toml
# minus russian-text-stresser (whose ancient pins break the resolver).
omegaconf==2.3.0
resampy==0.4.3
librosa
s3tokenizer
diffusers
resemble-perth==1.0.1
conformer
safetensors
spacy-pkuseg
pykakasi==2.3.0
backend/python/chatterbox/requirements.txt
+6
View File
@@ -0,0 +1,6 @@
grpcio==1.71.0
protobuf
certifi
packaging
setuptools
poetry
backend/python/chatterbox/run.sh
Executable
+9
View 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 $@
backend/python/chatterbox/test.py
+82
View File
@@ -0,0 +1,82 @@
"""
A test script to test the gRPC service
"""
import unittest
import subprocess
import time
import backend_pb2
import backend_pb2_grpc
import grpc
class TestBackendServicer(unittest.TestCase):
"""
TestBackendServicer is the class that tests the gRPC service
"""
def setUp(self):
"""
This method sets up the gRPC service by starting the server
"""
self.service = subprocess.Popen(["python3", "backend.py", "--addr", "localhost:50051"])
time.sleep(30)
def tearDown(self) -> None:
"""
This method tears down the gRPC service by terminating the server
"""
self.service.terminate()
self.service.wait()
def test_server_startup(self):
"""
This method tests if the server starts up successfully
"""
try:
self.setUp()
with grpc.insecure_channel("localhost:50051") as channel:
stub = backend_pb2_grpc.BackendStub(channel)
response = stub.Health(backend_pb2.HealthMessage())
self.assertEqual(response.message, b'OK')
except Exception as err:
print(err)
self.fail("Server failed to start")
finally:
self.tearDown()
def test_load_model(self):
"""
This method tests if the model is loaded successfully
"""
try:
self.setUp()
with grpc.insecure_channel("localhost:50051") as channel:
stub = backend_pb2_grpc.BackendStub(channel)
response = stub.LoadModel(backend_pb2.ModelOptions())
print(response)
self.assertTrue(response.success)
self.assertEqual(response.message, "Model loaded successfully")
except Exception as err:
print(err)
self.fail("LoadModel service failed")
finally:
self.tearDown()
def test_tts(self):
"""
This method tests if the embeddings are generated successfully
"""
try:
self.setUp()
with grpc.insecure_channel("localhost:50051") as channel:
stub = backend_pb2_grpc.BackendStub(channel)
response = stub.LoadModel(backend_pb2.ModelOptions())
self.assertTrue(response.success)
tts_request = backend_pb2.TTSRequest(text="80s TV news production music hit for tonight's biggest story")
tts_response = stub.TTS(tts_request)
self.assertIsNotNone(tts_response)
except Exception as err:
print(err)
self.fail("TTS service failed")
finally:
self.tearDown()
backend/python/chatterbox/test.sh
Executable
+11
View 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
backend/python/common/grpc_auth.py
+87
View File
@@ -0,0 +1,87 @@
"""Shared gRPC bearer token authentication interceptor for LocalAI Python backends.
When the environment variable LOCALAI_GRPC_AUTH_TOKEN is set, requests without
a valid Bearer token in the 'authorization' metadata header are rejected with
UNAUTHENTICATED. When the variable is empty or unset, no authentication is
performed (backward compatible).
"""
import hmac
import os
import grpc
from parent_watch import start_parent_death_watcher
class _AbortHandler(grpc.RpcMethodHandler):
"""A method handler that immediately aborts with UNAUTHENTICATED."""
def __init__(self):
self.request_streaming = False
self.response_streaming = False
self.request_deserializer = None
self.response_serializer = None
self.unary_unary = self._abort
self.unary_stream = None
self.stream_unary = None
self.stream_stream = None
@staticmethod
def _abort(request, context):
context.abort(grpc.StatusCode.UNAUTHENTICATED, "invalid token")
class TokenAuthInterceptor(grpc.ServerInterceptor):
"""Sync gRPC server interceptor that validates a bearer token."""
def __init__(self, token: str):
self._token = token
self._abort_handler = _AbortHandler()
def intercept_service(self, continuation, handler_call_details):
metadata = dict(handler_call_details.invocation_metadata)
auth = metadata.get("authorization", "")
expected = "Bearer " + self._token
if not hmac.compare_digest(auth, expected):
return self._abort_handler
return continuation(handler_call_details)
class AsyncTokenAuthInterceptor(grpc.aio.ServerInterceptor):
"""Async gRPC server interceptor that validates a bearer token."""
def __init__(self, token: str):
self._token = token
async def intercept_service(self, continuation, handler_call_details):
metadata = dict(handler_call_details.invocation_metadata)
auth = metadata.get("authorization", "")
expected = "Bearer " + self._token
if not hmac.compare_digest(auth, expected):
return _AbortHandler()
return await continuation(handler_call_details)
def get_auth_interceptors(*, aio: bool = False):
"""Return a list of gRPC interceptors for bearer token auth.
Args:
aio: If True, return async-compatible interceptors for grpc.aio.server().
If False (default), return sync interceptors for grpc.server().
Returns an empty list when LOCALAI_GRPC_AUTH_TOKEN is not set.
"""
# Arm the best-effort parent-death backstop here: this is the single helper
# every LocalAI Python backend invokes exactly once while building its gRPC
# server (mirroring how the Go watcher arms in pkg/grpc's shared serve path).
# start_parent_death_watcher() is idempotent and a no-op when disabled or on
# unsupported platforms — see parent_watch.py.
start_parent_death_watcher()
token = os.environ.get("LOCALAI_GRPC_AUTH_TOKEN", "")
if not token:
return []
if aio:
return [AsyncTokenAuthInterceptor(token)]
return [TokenAuthInterceptor(token)]
backend/python/common/libbackend.sh
+578
View File
@@ -0,0 +1,578 @@
#!/usr/bin/env bash
set -euo pipefail
#
# use the library by adding the following line to a script:
# source $(dirname $0)/../common/libbackend.sh
#
# If you want to limit what targets a backend can be used on, set the variable LIMIT_TARGETS to a
# space separated list of valid targets BEFORE sourcing the library, for example to only allow a backend
# to be used on CUDA and CPU backends:
#
# LIMIT_TARGETS="cublas cpu"
# source $(dirname $0)/../common/libbackend.sh
#
# You can use any valid BUILD_TYPE or BUILD_PROFILE, if you need to limit a backend to CUDA 12 only:
#
# LIMIT_TARGETS="cublas12"
# source $(dirname $0)/../common/libbackend.sh
#
# You can switch between uv (conda-like) and pip installation methods by setting USE_PIP:
# USE_PIP=true source $(dirname $0)/../common/libbackend.sh
#
# ===================== user-configurable defaults =====================
PYTHON_VERSION="${PYTHON_VERSION:-3.10}" # e.g. 3.10 / 3.11 / 3.12 / 3.13
PYTHON_PATCH="${PYTHON_PATCH:-18}" # e.g. 18 -> 3.10.18 ; 13 -> 3.11.13
PY_STANDALONE_TAG="${PY_STANDALONE_TAG:-20250818}" # release tag date
# Enable/disable bundling of a portable Python build
PORTABLE_PYTHON="${PORTABLE_PYTHON:-false}"
# If you want to fully pin the filename (including tuned CPU targets), set:
# PORTABLE_PY_FILENAME="cpython-3.10.18+20250818-x86_64_v3-unknown-linux-gnu-install_only.tar.gz"
: "${PORTABLE_PY_FILENAME:=}"
: "${PORTABLE_PY_SHA256:=}" # optional; if set we verify the download
# =====================================================================
# Default to uv if USE_PIP is not set
if [ "x${USE_PIP:-}" == "x" ]; then
USE_PIP=false
fi
# ----------------------- helpers -----------------------
function _is_musl() {
# detect musl (Alpine, etc)
if command -v ldd >/dev/null 2>&1; then
ldd --version 2>&1 | grep -qi musl && return 0
fi
# busybox-ish fallback
if command -v getconf >/dev/null 2>&1; then
getconf GNU_LIBC_VERSION >/dev/null 2>&1 || return 0
fi
return 1
}
function _triple() {
local os="" arch="" libc="gnu"
case "$(uname -s)" in
Linux*) os="unknown-linux" ;;
Darwin*) os="apple-darwin" ;;
MINGW*|MSYS*|CYGWIN*) os="pc-windows-msvc" ;; # best-effort for Git Bash
*) echo "Unsupported OS $(uname -s)"; exit 1;;
esac
case "$(uname -m)" in
x86_64) arch="x86_64" ;;
aarch64|arm64) arch="aarch64" ;;
armv7l) arch="armv7" ;;
i686|i386) arch="i686" ;;
ppc64le) arch="ppc64le" ;;
s390x) arch="s390x" ;;
riscv64) arch="riscv64" ;;
*) echo "Unsupported arch $(uname -m)"; exit 1;;
esac
if [[ "$os" == "unknown-linux" ]]; then
if _is_musl; then
libc="musl"
else
libc="gnu"
fi
echo "${arch}-${os}-${libc}"
else
echo "${arch}-${os}"
fi
}
function _portable_dir() {
echo "${EDIR}/python"
}
function _portable_bin() {
# python-build-standalone puts python in ./bin
echo "$(_portable_dir)/bin"
}
function _portable_python() {
if [ -x "$(_portable_bin)/python3" ]; then
echo "$(_portable_bin)/python3"
else
echo "$(_portable_bin)/python"
fi
}
# macOS loader env for the portable CPython
_macosPortableEnv() {
if [ "$(uname -s)" = "Darwin" ]; then
export DYLD_LIBRARY_PATH="$(_portable_dir)/lib${DYLD_LIBRARY_PATH:+:${DYLD_LIBRARY_PATH}}"
export DYLD_FALLBACK_LIBRARY_PATH="$(_portable_dir)/lib${DYLD_FALLBACK_LIBRARY_PATH:+:${DYLD_FALLBACK_LIBRARY_PATH}}"
fi
}
# Good hygiene on macOS for downloaded/extracted trees
_unquarantinePortablePython() {
if [ "$(uname -s)" = "Darwin" ]; then
command -v xattr >/dev/null 2>&1 && xattr -dr com.apple.quarantine "$(_portable_dir)" || true
fi
}
# ------------------ ### PORTABLE PYTHON ------------------
function ensurePortablePython() {
local pdir="$(_portable_dir)"
local pbin="$(_portable_bin)"
local pyexe
if [ -x "${pbin}/python3" ] || [ -x "${pbin}/python" ]; then
_macosPortableEnv
return 0
fi
mkdir -p "${pdir}"
local triple="$(_triple)"
local full_ver="${PYTHON_VERSION}.${PYTHON_PATCH}"
local fn=""
if [ -n "${PORTABLE_PY_FILENAME}" ]; then
fn="${PORTABLE_PY_FILENAME}"
else
# generic asset name: cpython-<full_ver>+<tag>-<triple>-install_only.tar.gz
fn="cpython-${full_ver}+${PY_STANDALONE_TAG}-${triple}-install_only.tar.gz"
fi
local url="https://github.com/astral-sh/python-build-standalone/releases/download/${PY_STANDALONE_TAG}/${fn}"
local tmp="${pdir}/${fn}"
echo "Downloading portable Python: ${fn}"
# curl with retries; fall back to wget if needed
if command -v curl >/dev/null 2>&1; then
curl -L --fail --retry 3 --retry-delay 1 -o "${tmp}" "${url}"
else
wget -O "${tmp}" "${url}"
fi
if [ -n "${PORTABLE_PY_SHA256}" ]; then
echo "${PORTABLE_PY_SHA256} ${tmp}" | sha256sum -c -
fi
echo "Extracting ${fn} -> ${pdir}"
# always a .tar.gz (we purposely choose install_only)
tar -xzf "${tmp}" -C "${pdir}"
rm -f "${tmp}"
# Some archives nest a directory; if so, flatten to ${pdir}
# Find the first dir with a 'bin/python*'
local inner
inner="$(find "${pdir}" -type f -path "*/bin/python*" -maxdepth 3 2>/dev/null | head -n1 || true)"
if [ -n "${inner}" ]; then
local inner_root
inner_root="$(dirname "$(dirname "${inner}")")" # .../bin -> root
if [ "${inner_root}" != "${pdir}" ]; then
# move contents up one level
shopt -s dotglob
mv "${inner_root}/"* "${pdir}/"
rm -rf "${inner_root}"
shopt -u dotglob
fi
fi
_unquarantinePortablePython
_macosPortableEnv
# Make sure it's runnable
pyexe="$(_portable_python)"
"${pyexe}" -V
}
# init handles the setup of the library
function init() {
BACKEND_NAME=${PWD##*/}
MY_DIR=$(realpath "$(dirname "$0")")
BUILD_PROFILE=$(getBuildProfile)
EDIR=${MY_DIR}
if [ "x${ENV_DIR:-}" != "x" ]; then
EDIR=${ENV_DIR}
fi
if [ ! -z "${LIMIT_TARGETS:-}" ]; then
isValidTarget=$(checkTargets ${LIMIT_TARGETS})
if [ ${isValidTarget} != true ]; then
echo "${BACKEND_NAME} can only be used on the following targets: ${LIMIT_TARGETS}"
exit 0
fi
fi
echo "Initializing libbackend for ${BACKEND_NAME}"
}
# getBuildProfile will inspect the system to determine which build profile is appropriate:
# returns one of the following:
# - cublas12
# - cublas13
# - hipblas
# - intel
function getBuildProfile() {
if [ x"${BUILD_TYPE:-}" == "xcublas" ] || [ x"${BUILD_TYPE:-}" == "xl4t" ]; then
if [ ! -z "${CUDA_MAJOR_VERSION:-}" ]; then
echo ${BUILD_TYPE}${CUDA_MAJOR_VERSION}
else
echo ${BUILD_TYPE}
fi
return 0
fi
if [ -d "/opt/intel" ]; then
echo "intel"
return 0
fi
if [ -n "${BUILD_TYPE:-}" ]; then
echo ${BUILD_TYPE}
return 0
fi
echo "cpu"
}
# Make the venv relocatable:
# - rewrite venv/bin/python{,3} to relative symlinks into $(_portable_dir)
# - normalize entrypoint shebangs to /usr/bin/env python3
# - optionally update pyvenv.cfg to point to the portable Python directory (only at runtime)
# Usage: _makeVenvPortable [--update-pyvenv-cfg]
_makeVenvPortable() {
local update_pyvenv_cfg=false
if [ "${1:-}" = "--update-pyvenv-cfg" ]; then
update_pyvenv_cfg=true
fi
local venv_dir="${EDIR}/venv"
local vbin="${venv_dir}/bin"
[ -d "${vbin}" ] || return 0
# 1) Replace python symlinks with relative ones to ../../python/bin/python3
# (venv/bin -> venv -> EDIR -> python/bin)
local rel_py='../../python/bin/python3'
for name in python3 python; do
if [ -e "${vbin}/${name}" ] || [ -L "${vbin}/${name}" ]; then
rm -f "${vbin}/${name}"
fi
done
ln -s "${rel_py}" "${vbin}/python3"
ln -s "python3" "${vbin}/python"
# 2) Update pyvenv.cfg to point to the portable Python directory (only at runtime)
# Use absolute path resolved at runtime so it works when the venv is copied
if [ "$update_pyvenv_cfg" = "true" ]; then
local pyvenv_cfg="${venv_dir}/pyvenv.cfg"
if [ -f "${pyvenv_cfg}" ]; then
local portable_dir="$(_portable_dir)"
# Resolve to absolute path - this ensures it works when the backend is copied
# Only resolve if the directory exists (it should if ensurePortablePython was called)
if [ -d "${portable_dir}" ]; then
portable_dir="$(cd "${portable_dir}" && pwd)"
else
# Fallback to relative path if directory doesn't exist yet
portable_dir="../python"
fi
local sed_i=(sed -i)
# macOS/BSD sed needs a backup suffix; GNU sed doesn't. Make it portable:
if sed --version >/dev/null 2>&1; then
sed_i=(sed -i)
else
sed_i=(sed -i '')
fi
# Update the home field in pyvenv.cfg
# Handle both absolute paths (starting with /) and relative paths
if grep -q "^home = " "${pyvenv_cfg}"; then
"${sed_i[@]}" "s|^home = .*|home = ${portable_dir}|" "${pyvenv_cfg}"
else
# If home field doesn't exist, add it
echo "home = ${portable_dir}" >> "${pyvenv_cfg}"
fi
fi
fi
# 3) Rewrite shebangs of entry points to use env, so the venv is relocatable
# Only touch text files that start with #! and reference the current venv.
local ve_abs="${vbin}/python"
local sed_i=(sed -i)
# macOS/BSD sed needs a backup suffix; GNU sed doesn't. Make it portable:
if sed --version >/dev/null 2>&1; then
sed_i=(sed -i)
else
sed_i=(sed -i '')
fi
for f in "${vbin}"/*; do
[ -f "$f" ] || continue
# Fast path: check first two bytes (#!)
head -c2 "$f" 2>/dev/null | grep -q '^#!' || continue
# Only rewrite if the shebang mentions the (absolute) venv python
if head -n1 "$f" | grep -Fq "${ve_abs}"; then
"${sed_i[@]}" '1s|^#!.*$|#!/usr/bin/env python3|' "$f"
chmod +x "$f" 2>/dev/null || true
fi
done
}
# Apply the venv to the current process: VIRTUAL_ENV, PATH, PYTHONHOME hygiene.
# Equivalent to the runtime portion of `source bin/activate`, but computed from
# $EDIR (resolved at runtime via realpath) instead of the path baked into
# bin/activate at venv-create time. `uv venv` (and `python -m venv`) both bake
# the create-time absolute path in, so sourcing activate on a relocated venv —
# e.g. one built at /vllm/venv inside a Docker stage and unpacked under
# /backends/cuda13-vllm-development/venv at runtime — silently prepends a
# stale, non-existent path to $PATH. Doing the setup ourselves sidesteps that;
# this is the same approach `uv run` takes internally.
_activateVenv() {
export VIRTUAL_ENV="${EDIR}/venv"
export PATH="${EDIR}/venv/bin:${PATH}"
unset PYTHONHOME
}
# ensureVenv makes sure that the venv for the backend both exists, and is activated.
#
# This function is idempotent, so you can call it as many times as you want and it will
# always result in an activated virtual environment
function ensureVenv() {
local interpreter=""
if [ "x${PORTABLE_PYTHON}" == "xtrue" ] || [ -e "$(_portable_python)" ]; then
echo "Using portable Python"
ensurePortablePython
interpreter="$(_portable_python)"
else
# Prefer system python${PYTHON_VERSION}, else python3, else fall back to bundled
if command -v python${PYTHON_VERSION} >/dev/null 2>&1; then
interpreter="python${PYTHON_VERSION}"
elif command -v python3 >/dev/null 2>&1; then
interpreter="python3"
else
echo "No suitable system Python found, bootstrapping portable build..."
ensurePortablePython
interpreter="$(_portable_python)"
fi
fi
if [ ! -d "${EDIR}/venv" ]; then
if [ "x${USE_PIP}" == "xtrue" ]; then
# --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"
_activateVenv
"${interpreter}" -m pip install --upgrade pip
else
if [ "x${PORTABLE_PYTHON}" == "xtrue" ]; then
uv venv --python "${interpreter}" "${EDIR}/venv"
else
uv venv --python "${PYTHON_VERSION}" "${EDIR}/venv"
fi
fi
if [ "x${PORTABLE_PYTHON}" == "xtrue" ]; then
# During install, only update symlinks and shebangs, not pyvenv.cfg
_makeVenvPortable
fi
fi
# We call it here to make sure that when we source a venv we can still use python as expected
if [ -x "$(_portable_python)" ]; then
_macosPortableEnv
fi
if [ "x${VIRTUAL_ENV:-}" != "x${EDIR}/venv" ]; then
_activateVenv
fi
}
function runProtogen() {
ensureVenv
# Match grpcio-tools to the grpcio already installed by the backend's
# requirements. grpcio and grpcio-tools are released in lockstep, and the
# protoc that grpcio-tools bundles stamps a Protobuf "gencode" version into
# backend_pb2.py. Left unpinned, `uv pip install grpcio-tools` pulls the
# newest release, whose newer gencode (e.g. 7.35.0) trips Protobuf's
# runtime >= gencode guarantee at import time when a backend caps the
# protobuf runtime lower (vLLM pins it to 6.33.6), crashing the backend with
# "grpc service not ready" before it ever loads a model. Pinning
# grpcio-tools to the installed grpcio version keeps the gencode in step with
# the runtime. Falls back to unpinned when grpcio isn't installed yet.
# See mudler/LocalAI#10718.
local grpcio_tools_spec="grpcio-tools"
local grpcio_version
grpcio_version="$(python -c 'import importlib.metadata as m; print(m.version("grpcio"))' 2>/dev/null || true)"
if [ -n "${grpcio_version}" ]; then
grpcio_tools_spec="grpcio-tools==${grpcio_version}"
fi
if [ "x${USE_PIP}" == "xtrue" ]; then
pip install "${grpcio_tools_spec}"
else
uv pip install "${grpcio_tools_spec}"
fi
pushd "${EDIR}" >/dev/null
# use the venv python (ensures correct interpreter & sys.path)
python -m grpc_tools.protoc -I../../ -I./ --python_out=. --grpc_python_out=. backend.proto
popd >/dev/null
}
# installRequirements looks for several requirements files and if they exist runs the install for them in order
#
# - requirements-install.txt
# - requirements.txt
# - requirements-${BUILD_TYPE}.txt
# - requirements-${BUILD_PROFILE}.txt
#
# BUILD_PROFILE is a more specific version of BUILD_TYPE, ex: cuda-12 or cuda-13
# it can also include some options that we do not have BUILD_TYPES for, ex: intel
#
# NOTE: for BUILD_PROFILE==intel, this function does NOT automatically use the Intel python package index.
# you may want to add the following line to a requirements-intel.txt if you use one:
#
# --index-url https://download.pytorch.org/whl/xpu
#
# If you need to add extra flags into the pip install command you can do so by setting the variable EXTRA_PIP_INSTALL_FLAGS
# before calling installRequirements. For example:
#
# source $(dirname $0)/../common/libbackend.sh
# EXTRA_PIP_INSTALL_FLAGS="--no-build-isolation"
# installRequirements
function installRequirements() {
ensureVenv
declare -a requirementFiles=(
"${EDIR}/requirements-install.txt"
"${EDIR}/requirements.txt"
"${EDIR}/requirements-${BUILD_TYPE:-}.txt"
)
if [ "x${BUILD_TYPE:-}" != "x${BUILD_PROFILE}" ]; then
requirementFiles+=("${EDIR}/requirements-${BUILD_PROFILE}.txt")
fi
if [ "x${BUILD_TYPE:-}" == "x" ]; then
requirementFiles+=("${EDIR}/requirements-cpu.txt")
fi
requirementFiles+=("${EDIR}/requirements-after.txt")
if [ "x${BUILD_TYPE:-}" != "x${BUILD_PROFILE}" ]; then
requirementFiles+=("${EDIR}/requirements-${BUILD_PROFILE}-after.txt")
fi
# This is needed to build wheels that e.g. depends on Python.h
if [ "x${PORTABLE_PYTHON}" == "xtrue" ]; then
export C_INCLUDE_PATH="${C_INCLUDE_PATH:-}:$(_portable_dir)/include/python${PYTHON_VERSION}"
fi
for reqFile in ${requirementFiles[@]}; do
if [ -f "${reqFile}" ]; then
echo "starting requirements install for ${reqFile}"
if [ "x${USE_PIP}" == "xtrue" ]; then
pip install ${EXTRA_PIP_INSTALL_FLAGS:-} --requirement "${reqFile}"
else
uv pip install ${EXTRA_PIP_INSTALL_FLAGS:-} --requirement "${reqFile}"
fi
echo "finished requirements install for ${reqFile}"
fi
done
runProtogen
}
# startBackend discovers and runs the backend GRPC server
#
# You can specify a specific backend file to execute by setting BACKEND_FILE before calling startBackend.
# example:
#
# source ../common/libbackend.sh
# BACKEND_FILE="${MY_DIR}/source/backend.py"
# startBackend $@
#
# valid filenames for autodiscovered backend servers are:
# - server.py
# - backend.py
# - ${BACKEND_NAME}.py
function startBackend() {
ensureVenv
# Update pyvenv.cfg before running to ensure paths are correct for current location
# This is critical when the backend position is dynamic (e.g., copied from container)
if [ "x${PORTABLE_PYTHON}" == "xtrue" ] || [ -x "$(_portable_python)" ]; then
_makeVenvPortable --update-pyvenv-cfg
fi
# Set up GPU library paths if a lib directory exists
# This allows backends to include their own GPU libraries (CUDA, ROCm, etc.)
if [ -d "${EDIR}/lib" ]; then
export LD_LIBRARY_PATH="${EDIR}/lib:${LD_LIBRARY_PATH:-}"
echo "Added ${EDIR}/lib to LD_LIBRARY_PATH for GPU libraries"
fi
if [ ! -z "${BACKEND_FILE:-}" ]; then
exec "${EDIR}/venv/bin/python" "${BACKEND_FILE}" "$@"
elif [ -e "${MY_DIR}/server.py" ]; then
exec "${EDIR}/venv/bin/python" "${MY_DIR}/server.py" "$@"
elif [ -e "${MY_DIR}/backend.py" ]; then
exec "${EDIR}/venv/bin/python" "${MY_DIR}/backend.py" "$@"
elif [ -e "${MY_DIR}/${BACKEND_NAME}.py" ]; then
exec "${EDIR}/venv/bin/python" "${MY_DIR}/${BACKEND_NAME}.py" "$@"
fi
}
# runUnittests discovers and runs python unittests
#
# You can specify a specific test file to use by setting TEST_FILE before calling runUnittests.
# example:
#
# source ../common/libbackend.sh
# TEST_FILE="${MY_DIR}/source/test.py"
# runUnittests $@
#
# be default a file named test.py in the backends directory will be used
function runUnittests() {
ensureVenv
if [ ! -z "${TEST_FILE:-}" ]; then
testDir=$(dirname "$(realpath "${TEST_FILE}")")
testFile=$(basename "${TEST_FILE}")
pushd "${testDir}" >/dev/null
python -m unittest "${testFile}"
popd >/dev/null
elif [ -f "${MY_DIR}/test.py" ]; then
pushd "${MY_DIR}" >/dev/null
python -m unittest test.py
popd >/dev/null
else
echo "no tests defined for ${BACKEND_NAME}"
fi
}
##################################################################################
# Below here are helper functions not intended to be used outside of the library #
##################################################################################
# checkTargets determines if the current BUILD_TYPE or BUILD_PROFILE is in a list of valid targets
function checkTargets() {
targets=$@
declare -a targets=($targets)
for target in ${targets[@]}; do
if [ "x${BUILD_TYPE:-}" == "x${target}" ]; then
echo true; return 0
fi
if [ "x${BUILD_PROFILE}" == "x${target}" ]; then
echo true; return 0
fi
done
echo false
}
init
backend/python/common/mlx_utils.py
+108
View File
@@ -0,0 +1,108 @@
"""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:
return "", text
if think_start not in text:
# Models like Qwen3.5 open assistant turns already INSIDE thinking, so
# the generated text carries only the closing tag. Everything before it
# is reasoning that would otherwise leak into the content.
if think_end and think_end in text:
head, _, tail = text.partition(think_end)
return head.strip(), tail.strip()
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
backend/python/common/mlx_utils_test.py
+75
View File
@@ -0,0 +1,75 @@
"""Unit tests for the mlx/mlx-vlm shared helpers (mlx_utils.py).
Run standalone (Python standard library only, no backend venv needed):
python3 -m unittest mlx_utils_test
These mirror the server-less helper tests in backend/python/mlx/test.py
(TestSharedHelpers), but live here so they run on any platform: the mlx
test module imports grpc/backend_pb2 at import time and needs the MLX venv,
whereas mlx_utils only needs the standard library.
"""
import types
import unittest
from mlx_utils import parse_tool_calls, split_reasoning
class TestSplitReasoning(unittest.TestCase):
def test_both_tags(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_implicit_opener_only_closing_tag(self):
# Qwen3.5 opens the assistant turn already inside thinking, so the
# output carries only the closing tag; everything before it is reasoning.
r, c = split_reasoning(
"The user is asking about the weather.\n</think>\n\nThe weather in Rome is sunny.",
"<think>",
"</think>",
)
self.assertEqual(r, "The user is asking about the weather.")
self.assertEqual(c, "The weather in Rome is sunny.")
def test_no_tags_at_all(self):
r, c = split_reasoning("just text", "<think>", "</think>")
self.assertEqual(r, "")
self.assertEqual(c, "just text")
def test_empty_think_end_and_no_opener_match(self):
# No think_end to anchor on, and the opener is absent → return unchanged.
r, c = split_reasoning("no opener here", "<think>", "")
self.assertEqual(r, "")
self.assertEqual(c, "no opener here")
def test_empty_text(self):
r, c = split_reasoning("", "<think>", "</think>")
self.assertEqual(r, "")
self.assertEqual(c, "")
class TestParseToolCalls(unittest.TestCase):
def test_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)
if __name__ == "__main__":
unittest.main()
backend/python/common/parent_watch.py
+149
View File
@@ -0,0 +1,149 @@
"""Parent-death watcher (best-effort backstop) for LocalAI Python backends.
LocalAI spawns each backend as a child process and, on a clean shutdown, tears
it down itself (SIGTERM -> grace -> SIGKILL). That graceful path only runs when
LocalAI receives a catchable signal and lives long enough to run its handlers.
If LocalAI is SIGKILLed (e.g. a supervising process's grace period elapses
first), that teardown never runs and this backend would be reparented to init
and linger, holding GPU/VRAM and its listen port.
The watcher here is a best-effort backstop for exactly that case: it does NOT
replace the graceful teardown, it only covers the "parent vanished without
cleaning up" path. It detects reparenting: when the process that spawned this
backend dies, the kernel reparents us to the nearest sub-reaper or to init
(PID 1), so os.getppid() stops matching the value captured at startup. This
getppid() approach is portable across Linux/macOS (unlike the Linux-only
PR_SET_PDEATHSIG), which is why it is used here, mirroring the Go backends'
pkg/grpc/parentwatch.go and the C++ backends' parent_watch.h. It is disabled on
Windows, which has no equivalent orphan-reparenting semantics.
Env vars (shared verbatim across the Go, C++ and Python backends):
LOCALAI_BACKEND_PARENT_WATCH enabled by default; a falsey value
("false"/"0"/"no"/"off", case-insensitive)
disables it.
LOCALAI_BACKEND_PARENT_WATCH_INTERVAL poll interval as a Go-style duration
string ("500ms", "2s", "1m") or a bare
number of seconds. Defaults to 2s.
"""
import os
import sys
import threading
ENV_PARENT_WATCH = "LOCALAI_BACKEND_PARENT_WATCH"
ENV_PARENT_WATCH_INTERVAL = "LOCALAI_BACKEND_PARENT_WATCH_INTERVAL"
_DEFAULT_INTERVAL_SECONDS = 2.0
# Guard so repeated calls (e.g. get_auth_interceptors invoked more than once)
# only ever arm a single watcher thread per process.
_started = False
_started_lock = threading.Lock()
def _enabled():
"""Report whether the watcher should run in this process."""
# Windows does not reparent orphans to a well-known init PID, so the
# getppid() heuristic used here doesn't apply there.
if os.name == "nt" or sys.platform.startswith("win"):
return False
val = os.environ.get(ENV_PARENT_WATCH, "").strip().lower()
if val in ("false", "0", "no", "off"):
return False
return True
def _interval_seconds():
"""Return the configured poll interval in seconds, or the default.
Accepts Go-style duration strings ("500ms", "2s", "1m") for cross-language
parity, or a bare number interpreted as seconds.
"""
raw = os.environ.get(ENV_PARENT_WATCH_INTERVAL, "").strip()
if not raw:
return _DEFAULT_INTERVAL_SECONDS
# Split numeric prefix from unit suffix.
i = 0
while i < len(raw) and (raw[i].isdigit() or raw[i] == "." or (i == 0 and raw[i] in "+-")):
i += 1
if i == 0:
return _DEFAULT_INTERVAL_SECONDS
try:
num = float(raw[:i])
except ValueError:
return _DEFAULT_INTERVAL_SECONDS
unit = raw[i:].lower()
if unit == "ms":
seconds = num / 1000.0
elif unit in ("s", ""):
seconds = num
elif unit == "m":
seconds = num * 60.0
else:
return _DEFAULT_INTERVAL_SECONDS
return seconds if seconds > 0 else _DEFAULT_INTERVAL_SECONDS
def _parent_died(orig_ppid):
"""Report whether this process has been reparented away from orig_ppid.
Reparenting is the standard POSIX signal that the original parent (here, the
LocalAI process that spawned this backend) has exited: the orphan is handed
to the nearest sub-reaper or to init (PID 1), so os.getppid() no longer
matches the value captured at startup.
"""
ppid = os.getppid()
return ppid != orig_ppid or ppid == 1
def _watch(orig_ppid, interval, on_death):
"""Poll until _parent_died reports the original parent is gone, then call
on_death. Blocks, so run it on its own (daemon) thread."""
import time
while True:
time.sleep(interval)
if _parent_died(orig_ppid):
on_death()
return
def start_parent_death_watcher():
"""Install the best-effort safety net described in this module's docstring.
No-op when disabled, on Windows, when already orphaned at startup
(os.getppid() <= 1), or if already started. This is a backstop alongside —
never a replacement for — LocalAI's graceful teardown.
"""
global _started
if not _enabled():
return
with _started_lock:
if _started:
return
orig_ppid = os.getppid()
# A parent of 1 (or less) at startup means we were already orphaned (or
# launched directly under init) — there is no original parent to watch.
if orig_ppid <= 1:
return
interval = _interval_seconds()
def on_death():
print(
"backend parent process (pid {}) exited without stopping this "
"backend; self-terminating to avoid orphaning".format(orig_ppid),
file=sys.stderr,
flush=True,
)
# Immediate, non-cleanup exit: this is a shutdown safety net and the
# normal graceful path is already gone.
os._exit(1)
thread = threading.Thread(
target=_watch,
args=(orig_ppid, interval, on_death),
name="parent-death-watcher",
daemon=True,
)
thread.start()
_started = True
backend/python/common/parent_watch_test.py
+150
View File
@@ -0,0 +1,150 @@
"""Unit tests for the parent-death watcher (parent_watch.py).
Run standalone (Python standard library only, no backend venv needed):
python3 -m unittest parent_watch_test
The core test (test_detects_reparent) builds a genuine two-level process tree
(test -> middle -> grandchild) with os.fork, lets the middle process die, and
asserts the grandchild's parent_watch._watch detects the reparenting and
self-terminates — mirroring the Go test in pkg/grpc/parentwatch_test.go and the
C++ test in backend/cpp/llama-cpp/parent_watch_test.cpp.
"""
import os
import sys
import tempfile
import threading
import time
import unittest
import parent_watch
class TestParentWatchEnvParsing(unittest.TestCase):
def setUp(self):
self._saved = {
k: os.environ.get(k)
for k in (parent_watch.ENV_PARENT_WATCH, parent_watch.ENV_PARENT_WATCH_INTERVAL)
}
for k in self._saved:
os.environ.pop(k, None)
def tearDown(self):
for k, v in self._saved.items():
if v is None:
os.environ.pop(k, None)
else:
os.environ[k] = v
def test_interval_default(self):
self.assertEqual(parent_watch._interval_seconds(), 2.0)
def test_interval_units(self):
cases = {"500ms": 0.5, "2s": 2.0, "1m": 60.0, "3": 3.0, "0.5s": 0.5}
for raw, expected in cases.items():
os.environ[parent_watch.ENV_PARENT_WATCH_INTERVAL] = raw
self.assertAlmostEqual(parent_watch._interval_seconds(), expected, msg=raw)
def test_interval_garbage_falls_back(self):
os.environ[parent_watch.ENV_PARENT_WATCH_INTERVAL] = "garbage"
self.assertEqual(parent_watch._interval_seconds(), 2.0)
@unittest.skipIf(os.name == "nt" or sys.platform.startswith("win"), "POSIX only")
def test_enabled_default(self):
self.assertTrue(parent_watch._enabled())
@unittest.skipIf(os.name == "nt" or sys.platform.startswith("win"), "POSIX only")
def test_disabled_by_falsey(self):
for val in ("false", "0", "no", "off", "OFF", " False "):
os.environ[parent_watch.ENV_PARENT_WATCH] = val
self.assertFalse(parent_watch._enabled(), msg=val)
@unittest.skipIf(os.name == "nt" or sys.platform.startswith("win"), "POSIX only")
def test_enabled_by_truthy(self):
for val in ("true", "1", "yes", "on"):
os.environ[parent_watch.ENV_PARENT_WATCH] = val
self.assertTrue(parent_watch._enabled(), msg=val)
@unittest.skipIf(os.name == "nt" or sys.platform.startswith("win"), "fork/reparent is POSIX only")
class TestParentWatchReparent(unittest.TestCase):
def _wait_for_file(self, path, timeout=10.0):
deadline = time.time() + timeout
while time.time() < deadline:
if os.path.exists(path):
return True
time.sleep(0.02)
return False
def test_detects_reparent(self):
tmpdir = tempfile.mkdtemp(prefix="parentwatch_test_")
ready_file = os.path.join(tmpdir, "ready")
exited_file = os.path.join(tmpdir, "exited")
middle = os.fork()
if middle == 0:
# ---- middle process ----
grandchild = os.fork()
if grandchild == 0:
# ---- grandchild process: arm the REAL watcher against middle ----
orig_ppid = os.getppid()
def on_death():
with open(exited_file, "w") as f:
f.write("1")
os._exit(7)
threading.Thread(
target=parent_watch._watch,
args=(orig_ppid, 0.05, on_death),
daemon=True,
).start()
# Safety valve: never linger if something goes wrong.
def bail():
time.sleep(30)
os._exit(2)
threading.Thread(target=bail, daemon=True).start()
# Signal readiness only after the watcher captured orig_ppid.
with open(ready_file, "w") as f:
f.write(str(os.getpid()))
while True:
time.sleep(1)
else:
# middle: wait until grandchild is ready, then exit to orphan it.
if not self._wait_for_file(ready_file):
os._exit(5)
os._exit(0)
# ---- test (top) process ----
os.waitpid(middle, 0) # reap middle only; grandchild is orphaned
self.assertTrue(os.path.exists(ready_file), "grandchild never signaled readiness")
self.assertTrue(
self._wait_for_file(exited_file),
"watcher did not detect parent death within timeout",
)
# Best-effort cleanup: kill the grandchild if it somehow survived.
try:
with open(ready_file) as f:
pid = int(f.read().strip())
if pid > 1:
os.kill(pid, 9)
except (OSError, ValueError):
pass
for p in (ready_file, exited_file):
try:
os.remove(p)
except OSError:
pass
try:
os.rmdir(tmpdir)
except OSError:
pass
if __name__ == "__main__":
unittest.main()
backend/python/common/python_utils.py
+76
View File
@@ -0,0 +1,76 @@
"""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:
tool_calls = json.loads(msg.tool_calls)
# Chat templates (e.g. Qwen) iterate function.arguments as a
# mapping, but the OpenAI wire format carries it as a JSON
# string — decode it back so the template's .items() works.
for tc in tool_calls:
fn = tc.get("function") if isinstance(tc, dict) else None
if isinstance(fn, dict) and isinstance(fn.get("arguments"), str):
try:
fn["arguments"] = json.loads(fn["arguments"])
except json.JSONDecodeError:
pass
d["tool_calls"] = tool_calls
except json.JSONDecodeError:
pass
result.append(d)
return result
backend/python/common/python_utils_test.py
+122
View File
@@ -0,0 +1,122 @@
"""Unit tests for the shared python backend helpers (python_utils.py).
Run standalone (Python standard library only, no backend venv needed):
python3 -m unittest python_utils_test
These mirror the server-less helper tests in backend/python/mlx/test.py
(TestSharedHelpers), but live here so they run on any platform: the mlx
test module imports grpc/backend_pb2 at import time and needs the MLX venv,
whereas python_utils has no third-party dependency. Proto Message objects
are faked with types.SimpleNamespace (real proto fields default to "").
"""
import json
import types
import unittest
from python_utils import messages_to_dicts, parse_options
def _msg(**fields):
"""Fake a proto Message: every unset field is the empty string, as protobuf."""
defaults = {
"role": "",
"content": "",
"name": "",
"tool_call_id": "",
"reasoning_content": "",
"tool_calls": "",
}
defaults.update(fields)
return types.SimpleNamespace(**defaults)
class TestParseOptions(unittest.TestCase):
def test_type_inference(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)
class TestMessagesToDicts(unittest.TestCase):
def test_basic_fields(self):
out = messages_to_dicts(
[
_msg(role="user", content="hi"),
_msg(role="tool", content="42", tool_call_id="call_1", name="f"),
]
)
self.assertEqual(out[0], {"role": "user", "content": "hi"})
self.assertEqual(out[1]["tool_call_id"], "call_1")
self.assertEqual(out[1]["name"], "f")
def test_tool_call_arguments_string_decoded_to_mapping(self):
# OpenAI wire format ships function.arguments as a JSON *string*; chat
# templates iterate it as a mapping, so it must come back as a dict.
out = messages_to_dicts(
[
_msg(
role="assistant",
tool_calls=json.dumps(
[
{
"id": "call_1",
"type": "function",
"function": {
"name": "get_weather",
"arguments": '{"location": "Rome"}',
},
}
]
),
)
]
)
args = out[0]["tool_calls"][0]["function"]["arguments"]
self.assertEqual(args, {"location": "Rome"})
self.assertEqual(dict(args.items()), {"location": "Rome"})
def test_tool_call_arguments_already_mapping_is_idempotent(self):
out = messages_to_dicts(
[
_msg(
role="assistant",
tool_calls=json.dumps(
[{"function": {"name": "f", "arguments": {"a": 1}}}]
),
)
]
)
self.assertEqual(out[0]["tool_calls"][0]["function"]["arguments"], {"a": 1})
def test_tool_call_arguments_invalid_json_left_as_string(self):
out = messages_to_dicts(
[
_msg(
role="assistant",
tool_calls=json.dumps(
[{"function": {"name": "f", "arguments": "not-json"}}]
),
)
]
)
self.assertEqual(out[0]["tool_calls"][0]["function"]["arguments"], "not-json")
def test_tool_call_without_function_key(self):
out = messages_to_dicts(
[_msg(role="assistant", tool_calls=json.dumps([{"id": "call_1"}]))]
)
self.assertEqual(out[0]["tool_calls"], [{"id": "call_1"}])
def test_tool_calls_invalid_json_dropped(self):
out = messages_to_dicts([_msg(role="assistant", tool_calls="{not json")])
self.assertNotIn("tool_calls", out[0])
if __name__ == "__main__":
unittest.main()
backend/python/common/template/Makefile
+13
View File
@@ -0,0 +1,13 @@
.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__
backend/python/common/template/backend.py
Executable
+4
View File
@@ -0,0 +1,4 @@
#!/usr/bin/env python3
import grpc
import backend_pb2
import backend_pb2_grpc
backend/python/common/template/install.sh
Executable
+19
View File
@@ -0,0 +1,19 @@
#!/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
# This is here because the Intel pip index is broken and returns 200 status codes for every package name, it just doesn't return any package links.
# This makes uv think that the package exists in the Intel pip index, and by default it stops looking at other pip indexes once it finds a match.
# We need uv to continue falling through to the pypi default index to find optimum[openvino] in the pypi index
# the --upgrade actually allows us to *downgrade* torch to the version provided in the Intel pip index
if [ "x${BUILD_PROFILE}" == "xintel" ]; then
EXTRA_PIP_INSTALL_FLAGS+=" --upgrade --index-strategy=unsafe-first-match"
fi
installRequirements
backend/python/common/template/protogen.sh
+11
View 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
runProtogen
backend/python/common/template/requirements-hipblas.txt
+2
View File
@@ -0,0 +1,2 @@
--extra-index-url https://download.pytorch.org/whl/rocm7.0
torch
backend/python/common/template/requirements-intel.txt
+4
View File
@@ -0,0 +1,4 @@
--extra-index-url https://download.pytorch.org/whl/xpu
torch==2.8.0
oneccl_bind_pt==2.8.0+xpu
optimum[openvino]
backend/python/common/template/requirements.txt
+3
View File
@@ -0,0 +1,3 @@
grpcio==1.80.0
protobuf
grpcio-tools
backend/python/common/template/run.sh
Executable
+9
View 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 $@
backend/python/common/template/test.sh
Executable
+11
View 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
backend/python/common/vllm_utils.py
+43
View File
@@ -0,0 +1,43 @@
"""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
__all__ = ["parse_options", "messages_to_dicts", "setup_parsers"]
def setup_parsers(opts):
"""Return ``(tool_parser_cls, reasoning_parser_cls)`` from an opts dict.
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
tool_parser_name = opts.get("tool_parser")
reasoning_parser_name = opts.get("reasoning_parser")
if tool_parser_name:
try:
from vllm.tool_parsers import ToolParserManager
tool_parser_cls = ToolParserManager.get_tool_parser(tool_parser_name)
print(f"[vllm_utils] Loaded tool_parser: {tool_parser_name}", file=sys.stderr)
except Exception as e:
print(f"[vllm_utils] Failed to load tool_parser {tool_parser_name}: {e}", file=sys.stderr)
if reasoning_parser_name:
try:
from vllm.reasoning import ReasoningParserManager
reasoning_parser_cls = ReasoningParserManager.get_reasoning_parser(reasoning_parser_name)
print(f"[vllm_utils] Loaded reasoning_parser: {reasoning_parser_name}", file=sys.stderr)
except Exception as e:
print(f"[vllm_utils] Failed to load reasoning_parser {reasoning_parser_name}: {e}", file=sys.stderr)
return tool_parser_cls, reasoning_parser_cls
backend/python/coqui/Makefile
+23
View File
@@ -0,0 +1,23 @@
.PHONY: coqui
coqui:
bash install.sh
.PHONY: run
run: coqui
@echo "Running coqui..."
bash run.sh
@echo "coqui run."
.PHONY: test
test: coqui
@echo "Testing coqui..."
bash test.sh
@echo "coqui tested."
.PHONY: protogen-clean
protogen-clean:
$(RM) backend_pb2_grpc.py backend_pb2.py
.PHONY: clean
clean: protogen-clean
rm -rf venv __pycache__
backend/python/coqui/README.md
+11
View File
@@ -0,0 +1,11 @@
# Creating a separate environment for coqui project
```
make coqui
```
# Testing the gRPC server
```
make test
```
backend/python/coqui/backend.py
+131
View File
@@ -0,0 +1,131 @@
#!/usr/bin/env python3
"""
This is an extra gRPC server of LocalAI for Coqui TTS
"""
from concurrent import futures
import time
import argparse
import signal
import sys
import os
import backend_pb2
import backend_pb2_grpc
import torch
from TTS.api import TTS
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
_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'))
COQUI_LANGUAGE = os.environ.get('COQUI_LANGUAGE', None)
# Implement the BackendServicer class with the service methods
class BackendServicer(backend_pb2_grpc.BackendServicer):
"""
BackendServicer is the class that implements the gRPC service
"""
def Health(self, request, context):
return backend_pb2.Reply(message=bytes("OK", 'utf-8'))
def LoadModel(self, request, context):
# Get device
# device = "cuda" if request.CUDA else "cpu"
if torch.cuda.is_available():
print("CUDA is available", file=sys.stderr)
device = "cuda"
else:
print("CUDA is not available", file=sys.stderr)
device = "cpu"
mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
if mps_available:
device = "mps"
if not torch.cuda.is_available() and request.CUDA:
return backend_pb2.Result(success=False, message="CUDA is not available")
self.AudioPath = None
# List available 🐸TTS models
print(TTS().list_models())
if os.path.isabs(request.AudioPath):
self.AudioPath = request.AudioPath
elif request.AudioPath and request.ModelFile != "" and not os.path.isabs(request.AudioPath):
# get base path of modelFile
modelFileBase = os.path.dirname(request.ModelFile)
# modify LoraAdapter to be relative to modelFileBase
self.AudioPath = os.path.join(modelFileBase, request.AudioPath)
try:
print("Preparing models, please wait", file=sys.stderr)
self.tts = TTS(request.Model).to(device)
except Exception as err:
return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")
# Implement your logic here for the LoadModel service
# Replace this with your desired response
return backend_pb2.Result(message="Model loaded successfully", success=True)
def TTS(self, request, context):
try:
# if model is multilingual add language from request or env as fallback
lang = request.language or COQUI_LANGUAGE
if lang == "":
lang = None
if self.tts.is_multi_lingual and lang is None:
return backend_pb2.Result(success=False, message=f"Model is multi-lingual, but no language was provided")
# if model is multi-speaker, use speaker_wav or the speaker_id from request.voice
if self.tts.is_multi_speaker and self.AudioPath is None and request.voice is None:
return backend_pb2.Result(success=False, message=f"Model is multi-speaker, but no speaker was provided")
if self.tts.is_multi_speaker and request.voice is not None:
self.tts.tts_to_file(text=request.text, speaker=request.voice, language=lang, file_path=request.dst)
else:
self.tts.tts_to_file(text=request.text, speaker_wav=self.AudioPath, language=lang, file_path=request.dst)
except Exception as err:
return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")
return backend_pb2.Result(success=True)
def serve(address):
server = grpc.server(futures.ThreadPoolExecutor(max_workers=MAX_WORKERS),
options=[
('grpc.max_message_length', 50 * 1024 * 1024), # 50MB
('grpc.max_send_message_length', 50 * 1024 * 1024), # 50MB
('grpc.max_receive_message_length', 50 * 1024 * 1024), # 50MB
],
interceptors=get_auth_interceptors(),
)
backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
server.add_insecure_port(address)
server.start()
print("Server started. Listening on: " + address, file=sys.stderr)
# Define the signal handler function
def signal_handler(sig, frame):
print("Received termination signal. Shutting down...")
server.stop(0)
sys.exit(0)
# Set the signal handlers for SIGINT and SIGTERM
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
try:
while True:
time.sleep(_ONE_DAY_IN_SECONDS)
except KeyboardInterrupt:
server.stop(0)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Run the gRPC server.")
parser.add_argument(
"--addr", default="localhost:50051", help="The address to bind the server to."
)
args = parser.parse_args()
serve(args.addr)
backend/python/coqui/install.sh
Executable
+19
View File
@@ -0,0 +1,19 @@
#!/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
# This is here because the Intel pip index is broken and returns 200 status codes for every package name, it just doesn't return any package links.
# This makes uv think that the package exists in the Intel pip index, and by default it stops looking at other pip indexes once it finds a match.
# We need uv to continue falling through to the pypi default index to find optimum[openvino] in the pypi index
# the --upgrade actually allows us to *downgrade* torch to the version provided in the Intel pip index
if [ "x${BUILD_PROFILE}" == "xintel" ]; then
EXTRA_PIP_INSTALL_FLAGS+=" --upgrade --index-strategy=unsafe-first-match"
fi
installRequirements
backend/python/coqui/requirements-cpu.txt
+6
View File
@@ -0,0 +1,6 @@
--extra-index-url https://download.pytorch.org/whl/cpu
transformers==4.48.3
accelerate
torch==2.4.1
torchaudio==2.4.1
coqui-tts
backend/python/coqui/requirements-cublas12.txt
+5
View File
@@ -0,0 +1,5 @@
torch==2.4.1
torchaudio==2.4.1
transformers==4.48.3
accelerate
coqui-tts
backend/python/coqui/requirements-hipblas.txt
+6
View File
@@ -0,0 +1,6 @@
--extra-index-url https://download.pytorch.org/whl/rocm7.0
torch==2.10.0+rocm7.0
torchaudio==2.10.0+rocm7.0
transformers==4.48.3
accelerate
coqui-tts
backend/python/coqui/requirements-intel.txt
+8
View File
@@ -0,0 +1,8 @@
--extra-index-url https://download.pytorch.org/whl/xpu
torch==2.8.0+xpu
torchaudio==2.8.0+xpu
optimum[openvino]
setuptools
transformers==4.48.3
accelerate
coqui-tts
backend/python/coqui/requirements-mps.txt
+4
View File
@@ -0,0 +1,4 @@
torch==2.7.1
transformers==4.48.3
accelerate
coqui-tts
backend/python/coqui/requirements.txt
+4
View File
@@ -0,0 +1,4 @@
grpcio==1.80.0
protobuf
certifi
packaging==26.2
backend/python/coqui/run.sh
Executable
+9
View 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 $@
backend/python/coqui/test.py
+82
View File
@@ -0,0 +1,82 @@
"""
A test script to test the gRPC service
"""
import unittest
import subprocess
import time
import backend_pb2
import backend_pb2_grpc
import grpc
class TestBackendServicer(unittest.TestCase):
"""
TestBackendServicer is the class that tests the gRPC service
"""
def setUp(self):
"""
This method sets up the gRPC service by starting the server
"""
self.service = subprocess.Popen(["python3", "backend.py", "--addr", "localhost:50051"])
time.sleep(30)
def tearDown(self) -> None:
"""
This method tears down the gRPC service by terminating the server
"""
self.service.terminate()
self.service.wait()
def test_server_startup(self):
"""
This method tests if the server starts up successfully
"""
try:
self.setUp()
with grpc.insecure_channel("localhost:50051") as channel:
stub = backend_pb2_grpc.BackendStub(channel)
response = stub.Health(backend_pb2.HealthMessage())
self.assertEqual(response.message, b'OK')
except Exception as err:
print(err)
self.fail("Server failed to start")
finally:
self.tearDown()
def test_load_model(self):
"""
This method tests if the model is loaded successfully
"""
try:
self.setUp()
with grpc.insecure_channel("localhost:50051") as channel:
stub = backend_pb2_grpc.BackendStub(channel)
response = stub.LoadModel(backend_pb2.ModelOptions(Model="tts_models/en/vctk/vits"))
print(response)
self.assertTrue(response.success)
self.assertEqual(response.message, "Model loaded successfully")
except Exception as err:
print(err)
self.fail("LoadModel service failed")
finally:
self.tearDown()
def test_tts(self):
"""
This method tests if the embeddings are generated successfully
"""
try:
self.setUp()
with grpc.insecure_channel("localhost:50051") as channel:
stub = backend_pb2_grpc.BackendStub(channel)
response = stub.LoadModel(backend_pb2.ModelOptions(Model="tts_models/en/vctk/vits"))
self.assertTrue(response.success)
tts_request = backend_pb2.TTSRequest(text="80s TV news production music hit for tonight's biggest story")
tts_response = stub.TTS(tts_request)
self.assertIsNotNone(tts_response)
except Exception as err:
print(err)
self.fail("TTS service failed")
finally:
self.tearDown()
backend/python/coqui/test.sh
Executable
+11
View 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
backend/python/diffusers/Makefile
+33
View File
@@ -0,0 +1,33 @@
export CONDA_ENV_PATH = "diffusers.yml"
ifeq ($(BUILD_TYPE), hipblas)
export CONDA_ENV_PATH = "diffusers-rocm.yml"
endif
# Intel GPU are supposed to have dependencies installed in the main python
# environment, so we skip conda installation for SYCL builds.
# https://github.com/intel/intel-extension-for-pytorch/issues/538
ifneq (,$(findstring sycl,$(BUILD_TYPE)))
export SKIP_CONDA=1
endif
.PHONY: diffusers
diffusers:
bash install.sh
.PHONY: run
run: diffusers
@echo "Running diffusers..."
bash run.sh
@echo "Diffusers run."
test: diffusers
bash test.sh
.PHONY: protogen-clean
protogen-clean:
$(RM) backend_pb2_grpc.py backend_pb2.py
.PHONY: clean
clean: protogen-clean
rm -rf venv __pycache__
backend/python/diffusers/README.md
+137
View File
@@ -0,0 +1,137 @@
# LocalAI Diffusers Backend
This backend provides gRPC access to Hugging Face diffusers pipelines with dynamic pipeline loading.
## Creating a separate environment for the diffusers project
```
make diffusers
```
## Dynamic Pipeline Loader
The diffusers backend includes a dynamic pipeline loader (`diffusers_dynamic_loader.py`) that automatically discovers and loads diffusers pipelines at runtime. This eliminates the need for per-pipeline conditional statements - new pipelines added to diffusers become available automatically without code changes.
### How It Works
1. **Pipeline Discovery**: On first use, the loader scans the `diffusers` package to find all classes that inherit from `DiffusionPipeline`.
2. **Registry Caching**: Discovery results are cached for the lifetime of the process to avoid repeated scanning.
3. **Task Aliases**: The loader automatically derives task aliases from class names (e.g., "text-to-image", "image-to-image", "inpainting") without hardcoding.
4. **Multiple Resolution Methods**: Pipelines can be resolved by:
- Exact class name (e.g., `StableDiffusionPipeline`)
- Task alias (e.g., `text-to-image`, `img2img`)
- Model ID (uses HuggingFace Hub to infer pipeline type)
### Usage Examples
```python
from diffusers_dynamic_loader import (
load_diffusers_pipeline,
get_available_pipelines,
get_available_tasks,
resolve_pipeline_class,
discover_diffusers_classes,
get_available_classes,
)
# List all available pipelines
pipelines = get_available_pipelines()
print(f"Available pipelines: {pipelines[:10]}...")
# List all task aliases
tasks = get_available_tasks()
print(f"Available tasks: {tasks}")
# Resolve a pipeline class by name
cls = resolve_pipeline_class(class_name="StableDiffusionPipeline")
# Resolve by task alias
cls = resolve_pipeline_class(task="stable-diffusion")
# Load and instantiate a pipeline
pipe = load_diffusers_pipeline(
class_name="StableDiffusionPipeline",
model_id="runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16
)
# Load from single file
pipe = load_diffusers_pipeline(
class_name="StableDiffusionPipeline",
model_id="/path/to/model.safetensors",
from_single_file=True,
torch_dtype=torch.float16
)
# Discover other diffusers classes (schedulers, models, etc.)
schedulers = discover_diffusers_classes("SchedulerMixin")
print(f"Available schedulers: {list(schedulers.keys())[:5]}...")
# Get list of available scheduler classes
scheduler_list = get_available_classes("SchedulerMixin")
```
### Generic Class Discovery
The dynamic loader can discover not just pipelines but any class type from diffusers:
```python
# Discover all scheduler classes
schedulers = discover_diffusers_classes("SchedulerMixin")
# Discover all model classes
models = discover_diffusers_classes("ModelMixin")
# Get a sorted list of available classes
scheduler_names = get_available_classes("SchedulerMixin")
```
### Special Pipeline Handling
Most pipelines are loaded dynamically through `load_diffusers_pipeline()`. Only pipelines requiring truly custom initialization logic are handled explicitly:
- `FluxTransformer2DModel`: Requires quantization and custom transformer loading (cannot use dynamic loader)
- `WanPipeline` / `WanImageToVideoPipeline`: Uses dynamic loader with special VAE (float32 dtype)
- `SanaPipeline`: Uses dynamic loader with post-load dtype conversion for VAE/text encoder
- `StableVideoDiffusionPipeline`: Uses dynamic loader with CPU offload handling
- `VideoDiffusionPipeline`: Alias for DiffusionPipeline with video flags
All other pipelines (StableDiffusionPipeline, StableDiffusionXLPipeline, FluxPipeline, etc.) are loaded purely through the dynamic loader.
### Error Handling
When a pipeline cannot be resolved, the loader provides helpful error messages listing available pipelines and tasks:
```
ValueError: Unknown pipeline class 'NonExistentPipeline'.
Available pipelines: AnimateDiffPipeline, AnimateDiffVideoToVideoPipeline, ...
```
## Environment Variables
| Variable | Default | Description |
|----------|---------|-------------|
| `COMPEL` | `0` | Enable Compel for prompt weighting |
| `SD_EMBED` | `0` | Enable sd_embed for prompt weighting |
| `XPU` | `0` | Enable Intel XPU support |
| `CLIPSKIP` | `1` | Enable CLIP skip support |
| `SAFETENSORS` | `1` | Use safetensors format |
| `CHUNK_SIZE` | `8` | Decode chunk size for video |
| `FPS` | `7` | Video frames per second |
| `DISABLE_CPU_OFFLOAD` | `0` | Disable CPU offload |
| `FRAMES` | `64` | Number of video frames |
| `BFL_REPO` | `ChuckMcSneed/FLUX.1-dev` | Flux base repo |
| `PYTHON_GRPC_MAX_WORKERS` | `1` | Max gRPC workers |
## Running Tests
```bash
./test.sh
```
The test suite includes:
- Unit tests for the dynamic loader (`test_dynamic_loader.py`)
- Integration tests for the gRPC backend (`test.py`)
backend/python/diffusers/backend.py
Executable
+1101
View File
File diff suppressed because it is too large Load Diff
backend/python/diffusers/diffusers_dynamic_loader.py
+538
View File
@@ -0,0 +1,538 @@
"""
Dynamic Diffusers Pipeline Loader
This module provides dynamic discovery and loading of diffusers pipelines at runtime,
eliminating the need for per-pipeline conditional statements. New pipelines added to
diffusers become available automatically without code changes.
The module also supports discovering other diffusers classes like schedulers, models,
and other components, making it a generic solution for dynamic class loading.
Usage:
from diffusers_dynamic_loader import load_diffusers_pipeline, get_available_pipelines
# Load by class name
pipe = load_diffusers_pipeline(class_name="StableDiffusionPipeline", model_id="...", torch_dtype=torch.float16)
# Load by task alias
pipe = load_diffusers_pipeline(task="text-to-image", model_id="...", torch_dtype=torch.float16)
# Load using model_id (infers from HuggingFace Hub if possible)
pipe = load_diffusers_pipeline(model_id="runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
# Get list of available pipelines
available = get_available_pipelines()
# Discover other diffusers classes (schedulers, models, etc.)
schedulers = discover_diffusers_classes("SchedulerMixin")
models = discover_diffusers_classes("ModelMixin")
"""
import importlib
import re
import sys
from typing import Any, Dict, List, Optional, Tuple, Type
# Global cache for discovered pipelines - computed once per process
_pipeline_registry: Optional[Dict[str, Type]] = None
_task_aliases: Optional[Dict[str, List[str]]] = None
# Global cache for other discovered class types
_class_registries: Dict[str, Dict[str, Type]] = {}
def _camel_to_kebab(name: str) -> str:
"""
Convert CamelCase to kebab-case.
Examples:
StableDiffusionPipeline -> stable-diffusion-pipeline
StableDiffusionXLImg2ImgPipeline -> stable-diffusion-xl-img-2-img-pipeline
"""
# Insert hyphen before uppercase letters (but not at the start)
s1 = re.sub('(.)([A-Z][a-z]+)', r'\1-\2', name)
# Insert hyphen before uppercase letters following lowercase letters or numbers
s2 = re.sub('([a-z0-9])([A-Z])', r'\1-\2', s1)
return s2.lower()
def _extract_task_keywords(class_name: str) -> List[str]:
"""
Extract task-related keywords from a pipeline class name.
This function derives useful task aliases from the class name without
hardcoding per-pipeline branches.
Returns a list of potential task aliases for this pipeline.
"""
aliases = []
name_lower = class_name.lower()
# Direct task mappings based on common patterns in class names
task_patterns = {
'text2image': ['text-to-image', 'txt2img', 'text2image'],
'texttoimage': ['text-to-image', 'txt2img', 'text2image'],
'txt2img': ['text-to-image', 'txt2img', 'text2image'],
'img2img': ['image-to-image', 'img2img', 'image2image'],
'image2image': ['image-to-image', 'img2img', 'image2image'],
'imagetoimage': ['image-to-image', 'img2img', 'image2image'],
'img2video': ['image-to-video', 'img2vid', 'img2video'],
'imagetovideo': ['image-to-video', 'img2vid', 'img2video'],
'text2video': ['text-to-video', 'txt2vid', 'text2video'],
'texttovideo': ['text-to-video', 'txt2vid', 'text2video'],
'inpaint': ['inpainting', 'inpaint'],
'depth2img': ['depth-to-image', 'depth2img'],
'depthtoimage': ['depth-to-image', 'depth2img'],
'controlnet': ['controlnet', 'control-net'],
'upscale': ['upscaling', 'upscale', 'super-resolution'],
'superresolution': ['upscaling', 'upscale', 'super-resolution'],
}
# Check for each pattern in the class name
for pattern, task_aliases in task_patterns.items():
if pattern in name_lower:
aliases.extend(task_aliases)
# Also detect general pipeline types from the class name structure
# E.g., StableDiffusionPipeline -> stable-diffusion, flux -> flux
# Remove "Pipeline" suffix and convert to kebab case
if class_name.endswith('Pipeline'):
base_name = class_name[:-8] # Remove "Pipeline"
kebab_name = _camel_to_kebab(base_name)
aliases.append(kebab_name)
# Extract model family name (e.g., "stable-diffusion" from "stable-diffusion-xl-img-2-img")
parts = kebab_name.split('-')
if len(parts) >= 2:
# Try the first two words as a family name
family = '-'.join(parts[:2])
if family not in aliases:
aliases.append(family)
# If no specific task pattern matched but class contains "Pipeline", add "text-to-image" as default
# since most diffusion pipelines support text-to-image generation
if 'text-to-image' not in aliases and 'image-to-image' not in aliases:
# Only add for pipelines that seem to be generation pipelines (not schedulers, etc.)
if 'pipeline' in name_lower and not any(x in name_lower for x in ['scheduler', 'processor', 'encoder']):
# Don't automatically add - let it be explicit
pass
return list(set(aliases)) # Remove duplicates
def discover_diffusers_classes(
base_class_name: str,
include_base: bool = True
) -> Dict[str, Type]:
"""
Discover all subclasses of a given base class from diffusers.
This function provides a generic way to discover any type of diffusers class,
not just pipelines. It can be used to discover schedulers, models, processors,
and other components.
Args:
base_class_name: Name of the base class to search for subclasses
(e.g., "DiffusionPipeline", "SchedulerMixin", "ModelMixin")
include_base: Whether to include the base class itself in results
Returns:
Dict mapping class names to class objects
Examples:
# Discover all pipeline classes
pipelines = discover_diffusers_classes("DiffusionPipeline")
# Discover all scheduler classes
schedulers = discover_diffusers_classes("SchedulerMixin")
# Discover all model classes
models = discover_diffusers_classes("ModelMixin")
# Discover AutoPipeline classes
auto_pipelines = discover_diffusers_classes("AutoPipelineForText2Image")
"""
global _class_registries
# Check cache first
if base_class_name in _class_registries:
return _class_registries[base_class_name]
import diffusers
# Try to get the base class from diffusers
base_class = None
try:
base_class = getattr(diffusers, base_class_name)
except AttributeError:
# Try to find in submodules
for submodule in ['schedulers', 'models', 'pipelines']:
try:
module = importlib.import_module(f'diffusers.{submodule}')
if hasattr(module, base_class_name):
base_class = getattr(module, base_class_name)
break
except (ImportError, ModuleNotFoundError):
continue
if base_class is None:
raise ValueError(f"Could not find base class '{base_class_name}' in diffusers")
registry: Dict[str, Type] = {}
# Include base class if requested
if include_base:
registry[base_class_name] = base_class
# Scan diffusers module for subclasses
for attr_name in dir(diffusers):
try:
attr = getattr(diffusers, attr_name)
if (isinstance(attr, type) and
issubclass(attr, base_class) and
(include_base or attr is not base_class)):
registry[attr_name] = attr
except (ImportError, AttributeError, TypeError, RuntimeError, ModuleNotFoundError):
continue
# Cache the results
_class_registries[base_class_name] = registry
return registry
def get_available_classes(base_class_name: str) -> List[str]:
"""
Get a sorted list of all discovered class names for a given base class.
Args:
base_class_name: Name of the base class (e.g., "SchedulerMixin")
Returns:
Sorted list of discovered class names
"""
return sorted(discover_diffusers_classes(base_class_name).keys())
def _discover_pipelines() -> Tuple[Dict[str, Type], Dict[str, List[str]]]:
"""
Discover all subclasses of DiffusionPipeline from diffusers.
This function uses the generic discover_diffusers_classes() internally
and adds pipeline-specific task alias generation. It also includes
AutoPipeline classes which are special utility classes for automatic
pipeline selection.
Returns:
A tuple of (pipeline_registry, task_aliases) where:
- pipeline_registry: Dict mapping class names to class objects
- task_aliases: Dict mapping task aliases to lists of class names
"""
# Use the generic discovery function
pipeline_registry = discover_diffusers_classes("DiffusionPipeline", include_base=True)
# Also add AutoPipeline classes - these are special utility classes that are
# NOT subclasses of DiffusionPipeline but are commonly used
import diffusers
auto_pipeline_classes = [
"AutoPipelineForText2Image",
"AutoPipelineForImage2Image",
"AutoPipelineForInpainting",
]
for cls_name in auto_pipeline_classes:
try:
cls = getattr(diffusers, cls_name)
if cls is not None:
pipeline_registry[cls_name] = cls
except AttributeError:
# Class not available in this version of diffusers
pass
# Generate task aliases for pipelines
task_aliases: Dict[str, List[str]] = {}
for attr_name in pipeline_registry:
if attr_name == "DiffusionPipeline":
continue # Skip base class for alias generation
aliases = _extract_task_keywords(attr_name)
for alias in aliases:
if alias not in task_aliases:
task_aliases[alias] = []
if attr_name not in task_aliases[alias]:
task_aliases[alias].append(attr_name)
return pipeline_registry, task_aliases
def get_pipeline_registry() -> Dict[str, Type]:
"""
Get the cached pipeline registry.
Returns a dictionary mapping pipeline class names to their class objects.
The registry is built on first access and cached for subsequent calls.
"""
global _pipeline_registry, _task_aliases
if _pipeline_registry is None:
_pipeline_registry, _task_aliases = _discover_pipelines()
return _pipeline_registry
def get_task_aliases() -> Dict[str, List[str]]:
"""
Get the cached task aliases dictionary.
Returns a dictionary mapping task aliases (e.g., "text-to-image") to
lists of pipeline class names that support that task.
"""
global _pipeline_registry, _task_aliases
if _task_aliases is None:
_pipeline_registry, _task_aliases = _discover_pipelines()
return _task_aliases
def get_available_pipelines() -> List[str]:
"""
Get a sorted list of all discovered pipeline class names.
Returns:
List of pipeline class names available for loading.
"""
return sorted(get_pipeline_registry().keys())
def get_available_tasks() -> List[str]:
"""
Get a sorted list of all available task aliases.
Returns:
List of task aliases (e.g., ["text-to-image", "image-to-image", ...])
"""
return sorted(get_task_aliases().keys())
def resolve_pipeline_class(
class_name: Optional[str] = None,
task: Optional[str] = None,
model_id: Optional[str] = None
) -> Type:
"""
Resolve a pipeline class from class_name, task, or model_id.
Priority:
1. If class_name is provided, look it up directly
2. If task is provided, resolve through task aliases
3. If model_id is provided, try to infer from HuggingFace Hub
Args:
class_name: Exact pipeline class name (e.g., "StableDiffusionPipeline")
task: Task alias (e.g., "text-to-image", "img2img")
model_id: HuggingFace model ID (e.g., "runwayml/stable-diffusion-v1-5")
Returns:
The resolved pipeline class.
Raises:
ValueError: If no pipeline could be resolved.
"""
registry = get_pipeline_registry()
aliases = get_task_aliases()
# 1. Direct class name lookup
if class_name:
if class_name in registry:
return registry[class_name]
# Try case-insensitive match
for name, cls in registry.items():
if name.lower() == class_name.lower():
return cls
raise ValueError(
f"Unknown pipeline class '{class_name}'. "
f"Available pipelines: {', '.join(sorted(registry.keys())[:20])}..."
)
# 2. Task alias lookup
if task:
task_lower = task.lower().replace('_', '-')
if task_lower in aliases:
# Return the first matching pipeline for this task
matching_classes = aliases[task_lower]
if matching_classes:
return registry[matching_classes[0]]
# Try partial matching
for alias, classes in aliases.items():
if task_lower in alias or alias in task_lower:
if classes:
return registry[classes[0]]
raise ValueError(
f"Unknown task '{task}'. "
f"Available tasks: {', '.join(sorted(aliases.keys())[:20])}..."
)
# 3. Try to infer from HuggingFace Hub
if model_id:
try:
from huggingface_hub import model_info
info = model_info(model_id)
# Check pipeline_tag
if hasattr(info, 'pipeline_tag') and info.pipeline_tag:
tag = info.pipeline_tag.lower().replace('_', '-')
if tag in aliases:
matching_classes = aliases[tag]
if matching_classes:
return registry[matching_classes[0]]
# Check model card for hints
if hasattr(info, 'cardData') and info.cardData:
card = info.cardData
if 'pipeline_tag' in card:
tag = card['pipeline_tag'].lower().replace('_', '-')
if tag in aliases:
matching_classes = aliases[tag]
if matching_classes:
return registry[matching_classes[0]]
except ImportError:
# huggingface_hub not available
pass
except (KeyError, AttributeError, ValueError, OSError):
# Model info lookup failed - common cases:
# - KeyError: Missing keys in model card
# - AttributeError: Missing attributes on model info
# - ValueError: Invalid model data
# - OSError: Network or file access issues
pass
# Fallback: use DiffusionPipeline.from_pretrained which auto-detects
# DiffusionPipeline is always added to registry in _discover_pipelines (line 132)
# but use .get() with import fallback for extra safety
from diffusers import DiffusionPipeline
return registry.get('DiffusionPipeline', DiffusionPipeline)
raise ValueError(
"Must provide at least one of: class_name, task, or model_id. "
f"Available pipelines: {', '.join(sorted(registry.keys())[:20])}... "
f"Available tasks: {', '.join(sorted(aliases.keys())[:20])}..."
)
def load_diffusers_pipeline(
class_name: Optional[str] = None,
task: Optional[str] = None,
model_id: Optional[str] = None,
from_single_file: bool = False,
**kwargs
) -> Any:
"""
Load a diffusers pipeline dynamically.
This function resolves the appropriate pipeline class based on the provided
parameters and instantiates it with the given kwargs.
Args:
class_name: Exact pipeline class name (e.g., "StableDiffusionPipeline")
task: Task alias (e.g., "text-to-image", "img2img")
model_id: HuggingFace model ID or local path
from_single_file: If True, use from_single_file() instead of from_pretrained()
**kwargs: Additional arguments passed to from_pretrained() or from_single_file()
Returns:
An instantiated pipeline object.
Raises:
ValueError: If no pipeline could be resolved.
Exception: If pipeline loading fails.
Examples:
# Load by class name
pipe = load_diffusers_pipeline(
class_name="StableDiffusionPipeline",
model_id="runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16
)
# Load by task
pipe = load_diffusers_pipeline(
task="text-to-image",
model_id="runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16
)
# Load from single file
pipe = load_diffusers_pipeline(
class_name="StableDiffusionPipeline",
model_id="/path/to/model.safetensors",
from_single_file=True,
torch_dtype=torch.float16
)
"""
# Resolve the pipeline class
pipeline_class = resolve_pipeline_class(
class_name=class_name,
task=task,
model_id=model_id
)
# If no model_id provided but we have a class, we can't load
if model_id is None:
raise ValueError("model_id is required to load a pipeline")
# Load the pipeline
try:
if from_single_file:
# Check if the class has from_single_file method
if hasattr(pipeline_class, 'from_single_file'):
return pipeline_class.from_single_file(model_id, **kwargs)
else:
raise ValueError(
f"Pipeline class {pipeline_class.__name__} does not support from_single_file(). "
f"Use from_pretrained() instead."
)
else:
return pipeline_class.from_pretrained(model_id, **kwargs)
except Exception as e:
# Provide helpful error message
available = get_available_pipelines()
raise RuntimeError(
f"Failed to load pipeline '{pipeline_class.__name__}' from '{model_id}': {e}\n"
f"Available pipelines: {', '.join(available[:20])}..."
) from e
def get_pipeline_info(class_name: str) -> Dict[str, Any]:
"""
Get information about a specific pipeline class.
Args:
class_name: The pipeline class name
Returns:
Dictionary with pipeline information including:
- name: Class name
- aliases: List of task aliases
- supports_single_file: Whether from_single_file() is available
- docstring: Class docstring (if available)
"""
registry = get_pipeline_registry()
aliases = get_task_aliases()
if class_name not in registry:
raise ValueError(f"Unknown pipeline: {class_name}")
cls = registry[class_name]
# Find all aliases for this pipeline
pipeline_aliases = []
for alias, classes in aliases.items():
if class_name in classes:
pipeline_aliases.append(alias)
return {
'name': class_name,
'aliases': pipeline_aliases,
'supports_single_file': hasattr(cls, 'from_single_file'),
'docstring': cls.__doc__[:200] if cls.__doc__ else None
}
backend/python/diffusers/install.sh
Executable
+30
View File
@@ -0,0 +1,30 @@
#!/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
# This is here because the Intel pip index is broken and returns 200 status codes for every package name, it just doesn't return any package links.
# This makes uv think that the package exists in the Intel pip index, and by default it stops looking at other pip indexes once it finds a match.
# We need uv to continue falling through to the pypi default index to find optimum[openvino] in the pypi index
# the --upgrade actually allows us to *downgrade* torch to the version provided in the Intel pip index
if [ "x${BUILD_PROFILE}" == "xintel" ]; then
EXTRA_PIP_INSTALL_FLAGS+=" --upgrade --index-strategy=unsafe-first-match"
fi
if [ "x${BUILD_PROFILE}" == "xl4t12" ]; then
USE_PIP=true
fi
# Use python 3.12 for l4t
if [ "x${BUILD_PROFILE}" == "xl4t13" ]; then
PYTHON_VERSION="3.12"
PYTHON_PATCH="12"
PY_STANDALONE_TAG="20251120"
fi
installRequirements
backend/python/diffusers/requirements-cpu.txt
+24
View File
@@ -0,0 +1,24 @@
--extra-index-url https://download.pytorch.org/whl/cpu
diffusers==0.38.0
opencv-python
transformers==4.57.6
torchvision==0.22.1
accelerate
git+https://github.com/xhinker/sd_embed
peft
sentencepiece
torch==2.7.1
optimum-quanto
ftfy
# diffusers and transformers are pinned together on purpose. transformers v5
# restructured CLIPTextModel and dropped the `.text_model` attribute, which
# breaks single-file Stable Diffusion loading on every released diffusers
# (<=0.38.0); only unreleased diffusers main supports transformers v5. Tracking
# main via git froze whichever broken pair existed at image-build time. Pin the
# last known-good released pair so builds are reproducible and can't drift into
# the broken window. See https://github.com/mudler/LocalAI/issues/9979
#
# compel is intentionally omitted: it pins transformers~=4.25, which conflicts
# with this pin and previously forced pip into multi-hour resolver backtracking
# storms in CI. backend.py imports it lazily and gates the COMPEL=1 env var on
# the import succeeding, so dropping it here is safe.
backend/python/diffusers/requirements-cublas12.txt
+24
View File
@@ -0,0 +1,24 @@
--extra-index-url https://download.pytorch.org/whl/cu121
diffusers==0.38.0
opencv-python
transformers==4.57.6
torchvision
accelerate
git+https://github.com/xhinker/sd_embed
peft
sentencepiece
torch
ftfy
optimum-quanto
# diffusers and transformers are pinned together on purpose. transformers v5
# restructured CLIPTextModel and dropped the `.text_model` attribute, which
# breaks single-file Stable Diffusion loading on every released diffusers
# (<=0.38.0); only unreleased diffusers main supports transformers v5. Tracking
# main via git froze whichever broken pair existed at image-build time. Pin the
# last known-good released pair so builds are reproducible and can't drift into
# the broken window. See https://github.com/mudler/LocalAI/issues/9979
#
# compel is intentionally omitted: it pins transformers~=4.25, which conflicts
# with this pin and previously forced pip into multi-hour resolver backtracking
# storms in CI. backend.py imports it lazily and gates the COMPEL=1 env var on
# the import succeeding, so dropping it here is safe.
backend/python/diffusers/requirements-cublas13.txt
+24
View File
@@ -0,0 +1,24 @@
--extra-index-url https://download.pytorch.org/whl/cu130
diffusers==0.38.0
opencv-python
transformers==4.57.6
torchvision
accelerate
git+https://github.com/xhinker/sd_embed
peft
sentencepiece
torch
ftfy
optimum-quanto
# diffusers and transformers are pinned together on purpose. transformers v5
# restructured CLIPTextModel and dropped the `.text_model` attribute, which
# breaks single-file Stable Diffusion loading on every released diffusers
# (<=0.38.0); only unreleased diffusers main supports transformers v5. Tracking
# main via git froze whichever broken pair existed at image-build time. Pin the
# last known-good released pair so builds are reproducible and can't drift into
# the broken window. See https://github.com/mudler/LocalAI/issues/9979
#
# compel is intentionally omitted: it pins transformers~=4.25, which conflicts
# with this pin and previously forced pip into multi-hour resolver backtracking
# storms in CI. backend.py imports it lazily and gates the COMPEL=1 env var on
# the import succeeding, so dropping it here is safe.
backend/python/diffusers/requirements-hipblas.txt
+23
View File
@@ -0,0 +1,23 @@
--extra-index-url https://download.pytorch.org/whl/rocm7.0
torch==2.10.0+rocm7.0
torchvision==0.25.0+rocm7.0
diffusers==0.38.0
opencv-python
transformers==4.57.6
accelerate
peft
sentencepiece
optimum-quanto
ftfy
# diffusers and transformers are pinned together on purpose. transformers v5
# restructured CLIPTextModel and dropped the `.text_model` attribute, which
# breaks single-file Stable Diffusion loading on every released diffusers
# (<=0.38.0); only unreleased diffusers main supports transformers v5. Tracking
# main via git froze whichever broken pair existed at image-build time. Pin the
# last known-good released pair so builds are reproducible and can't drift into
# the broken window. See https://github.com/mudler/LocalAI/issues/9979
#
# compel is intentionally omitted: it pins transformers~=4.25, which conflicts
# with this pin and previously forced pip into multi-hour resolver backtracking
# storms in CI. backend.py imports it lazily and gates the COMPEL=1 env var on
# the import succeeding, so dropping it here is safe.
backend/python/diffusers/requirements-intel.txt
+26
View File
@@ -0,0 +1,26 @@
--extra-index-url https://download.pytorch.org/whl/xpu
torch
torchvision
optimum[openvino]
setuptools
diffusers==0.38.0
opencv-python
transformers==4.57.6
accelerate
git+https://github.com/xhinker/sd_embed
peft
sentencepiece
optimum-quanto
ftfy
# diffusers and transformers are pinned together on purpose. transformers v5
# restructured CLIPTextModel and dropped the `.text_model` attribute, which
# breaks single-file Stable Diffusion loading on every released diffusers
# (<=0.38.0); only unreleased diffusers main supports transformers v5. Tracking
# main via git froze whichever broken pair existed at image-build time. Pin the
# last known-good released pair so builds are reproducible and can't drift into
# the broken window. See https://github.com/mudler/LocalAI/issues/9979
#
# compel is intentionally omitted: it pins transformers~=4.25, which conflicts
# with this pin and previously forced pip into multi-hour resolver backtracking
# storms in CI. backend.py imports it lazily and gates the COMPEL=1 env var on
# the import succeeding, so dropping it here is safe.
backend/python/diffusers/requirements-l4t12.txt
+23
View File
@@ -0,0 +1,23 @@
--extra-index-url https://pypi.jetson-ai-lab.io/jp6/cu129/
torch
diffusers==0.38.0
transformers==4.57.6
accelerate
peft
optimum-quanto
numpy<2
sentencepiece
torchvision
ftfy
# diffusers and transformers are pinned together on purpose. transformers v5
# restructured CLIPTextModel and dropped the `.text_model` attribute, which
# breaks single-file Stable Diffusion loading on every released diffusers
# (<=0.38.0); only unreleased diffusers main supports transformers v5. Tracking
# main via git froze whichever broken pair existed at image-build time. Pin the
# last known-good released pair so builds are reproducible and can't drift into
# the broken window. See https://github.com/mudler/LocalAI/issues/9979
#
# compel is intentionally omitted: it pins transformers~=4.25, which conflicts
# with this pin and previously forced pip into multi-hour resolver backtracking
# storms in CI. backend.py imports it lazily and gates the COMPEL=1 env var on
# the import succeeding, so dropping it here is safe.
backend/python/diffusers/requirements-l4t13.txt
+24
View File
@@ -0,0 +1,24 @@
--extra-index-url https://download.pytorch.org/whl/cu130
torch
diffusers==0.38.0
transformers==4.57.6
accelerate
peft
optimum-quanto
numpy<2
sentencepiece
torchvision
ftfy
chardet
# diffusers and transformers are pinned together on purpose. transformers v5
# restructured CLIPTextModel and dropped the `.text_model` attribute, which
# breaks single-file Stable Diffusion loading on every released diffusers
# (<=0.38.0); only unreleased diffusers main supports transformers v5. Tracking
# main via git froze whichever broken pair existed at image-build time. Pin the
# last known-good released pair so builds are reproducible and can't drift into
# the broken window. See https://github.com/mudler/LocalAI/issues/9979
#
# compel is intentionally omitted: it pins transformers~=4.25, which conflicts
# with this pin and previously forced pip into multi-hour resolver backtracking
# storms in CI. backend.py imports it lazily and gates the COMPEL=1 env var on
# the import succeeding, so dropping it here is safe.
backend/python/diffusers/requirements-mps.txt
+22
View File
@@ -0,0 +1,22 @@
torch==2.7.1
torchvision==0.22.1
diffusers==0.38.0
opencv-python
transformers==4.57.6
accelerate
peft
sentencepiece
optimum-quanto
ftfy
# diffusers and transformers are pinned together on purpose. transformers v5
# restructured CLIPTextModel and dropped the `.text_model` attribute, which
# breaks single-file Stable Diffusion loading on every released diffusers
# (<=0.38.0); only unreleased diffusers main supports transformers v5. Tracking
# main via git froze whichever broken pair existed at image-build time. Pin the
# last known-good released pair so builds are reproducible and can't drift into
# the broken window. See https://github.com/mudler/LocalAI/issues/9979
#
# compel is intentionally omitted: it pins transformers~=4.25, which conflicts
# with this pin and previously forced pip into multi-hour resolver backtracking
# storms in CI. backend.py imports it lazily and gates the COMPEL=1 env var on
# the import succeeding, so dropping it here is safe.
backend/python/diffusers/requirements.txt
+6
View File
@@ -0,0 +1,6 @@
setuptools
grpcio==1.76.0
pillow
protobuf
certifi
av
backend/python/diffusers/run.sh
Executable
+17
View File
@@ -0,0 +1,17 @@
#!/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
if [ -d "/opt/intel" ]; then
# Assumes we are using the Intel oneAPI container image
# https://github.com/intel/intel-extension-for-pytorch/issues/538
export XPU=1
fi
export PYTORCH_ENABLE_MPS_FALLBACK=1
startBackend $@
backend/python/diffusers/test.py
+375
View File
@@ -0,0 +1,375 @@
"""
A test script to test the gRPC service and dynamic loader
"""
import unittest
import subprocess
import time
from unittest.mock import patch, MagicMock
# Import dynamic loader for testing (these don't need gRPC)
import diffusers_dynamic_loader as loader
from diffusers import DiffusionPipeline, StableDiffusionPipeline
# Try to import gRPC modules - may not be available during unit testing
try:
import grpc
import backend_pb2
import backend_pb2_grpc
GRPC_AVAILABLE = True
except ImportError:
GRPC_AVAILABLE = False
@unittest.skipUnless(GRPC_AVAILABLE, "gRPC modules not available")
class TestBackendServicer(unittest.TestCase):
"""
TestBackendServicer is the class that tests the gRPC service
"""
def setUp(self):
"""
This method sets up the gRPC service by starting the server
"""
self.service = subprocess.Popen(["python3", "backend.py", "--addr", "localhost:50051"])
def tearDown(self) -> None:
"""
This method tears down the gRPC service by terminating the server
"""
self.service.kill()
self.service.wait()
def test_server_startup(self):
"""
This method tests if the server starts up successfully
"""
time.sleep(20)
try:
self.setUp()
with grpc.insecure_channel("localhost:50051") as channel:
stub = backend_pb2_grpc.BackendStub(channel)
response = stub.Health(backend_pb2.HealthMessage())
self.assertEqual(response.message, b'OK')
except Exception as err:
print(err)
self.fail("Server failed to start")
finally:
self.tearDown()
def test_load_model(self):
"""
This method tests if the model is loaded successfully
"""
time.sleep(20)
try:
self.setUp()
with grpc.insecure_channel("localhost:50051") as channel:
stub = backend_pb2_grpc.BackendStub(channel)
response = stub.LoadModel(backend_pb2.ModelOptions(Model="Lykon/dreamshaper-8"))
self.assertTrue(response.success)
self.assertEqual(response.message, "Model loaded successfully")
except Exception as err:
print(err)
self.fail("LoadModel service failed")
finally:
self.tearDown()
def test(self):
"""
This method tests if the backend can generate images
"""
time.sleep(20)
try:
self.setUp()
with grpc.insecure_channel("localhost:50051") as channel:
stub = backend_pb2_grpc.BackendStub(channel)
response = stub.LoadModel(backend_pb2.ModelOptions(Model="Lykon/dreamshaper-8"))
print(response.message)
self.assertTrue(response.success)
image_req = backend_pb2.GenerateImageRequest(positive_prompt="cat", width=16,height=16, dst="test.jpg")
re = stub.GenerateImage(image_req)
self.assertTrue(re.success)
except Exception as err:
print(err)
self.fail("Image gen service failed")
finally:
self.tearDown()
class TestDiffusersDynamicLoader(unittest.TestCase):
"""Test cases for the diffusers dynamic loader functionality."""
@classmethod
def setUpClass(cls):
"""Set up test fixtures - clear caches to ensure fresh discovery."""
# Reset the caches to ensure fresh discovery
loader._pipeline_registry = None
loader._task_aliases = None
def test_camel_to_kebab_conversion(self):
"""Test CamelCase to kebab-case conversion."""
test_cases = [
("StableDiffusionPipeline", "stable-diffusion-pipeline"),
("StableDiffusionXLPipeline", "stable-diffusion-xl-pipeline"),
("FluxPipeline", "flux-pipeline"),
("DiffusionPipeline", "diffusion-pipeline"),
]
for input_val, expected in test_cases:
with self.subTest(input=input_val):
result = loader._camel_to_kebab(input_val)
self.assertEqual(result, expected)
def test_extract_task_keywords(self):
"""Test task keyword extraction from class names."""
# Test text-to-image detection
aliases = loader._extract_task_keywords("StableDiffusionPipeline")
self.assertIn("stable-diffusion", aliases)
# Test img2img detection
aliases = loader._extract_task_keywords("StableDiffusionImg2ImgPipeline")
self.assertIn("image-to-image", aliases)
self.assertIn("img2img", aliases)
# Test inpainting detection
aliases = loader._extract_task_keywords("StableDiffusionInpaintPipeline")
self.assertIn("inpainting", aliases)
self.assertIn("inpaint", aliases)
# Test depth2img detection
aliases = loader._extract_task_keywords("StableDiffusionDepth2ImgPipeline")
self.assertIn("depth-to-image", aliases)
def test_discover_pipelines_finds_known_classes(self):
"""Test that pipeline discovery finds at least one known pipeline class."""
registry = loader.get_pipeline_registry()
# Check that the registry is not empty
self.assertGreater(len(registry), 0, "Pipeline registry should not be empty")
# Check for known pipeline classes
known_pipelines = [
"StableDiffusionPipeline",
"DiffusionPipeline",
]
for pipeline_name in known_pipelines:
with self.subTest(pipeline=pipeline_name):
self.assertIn(
pipeline_name,
registry,
f"Expected to find {pipeline_name} in registry"
)
def test_discover_pipelines_caches_results(self):
"""Test that pipeline discovery results are cached."""
# Get registry twice
registry1 = loader.get_pipeline_registry()
registry2 = loader.get_pipeline_registry()
# Should be the same object (cached)
self.assertIs(registry1, registry2, "Registry should be cached")
def test_get_available_pipelines(self):
"""Test getting list of available pipelines."""
available = loader.get_available_pipelines()
# Should return a list
self.assertIsInstance(available, list)
# Should contain known pipelines
self.assertIn("StableDiffusionPipeline", available)
self.assertIn("DiffusionPipeline", available)
# Should be sorted
self.assertEqual(available, sorted(available))
def test_get_available_tasks(self):
"""Test getting list of available task aliases."""
tasks = loader.get_available_tasks()
# Should return a list
self.assertIsInstance(tasks, list)
# Should be sorted
self.assertEqual(tasks, sorted(tasks))
def test_resolve_pipeline_class_by_name(self):
"""Test resolving pipeline class by exact name."""
cls = loader.resolve_pipeline_class(class_name="StableDiffusionPipeline")
self.assertEqual(cls, StableDiffusionPipeline)
def test_resolve_pipeline_class_by_name_case_insensitive(self):
"""Test that class name resolution is case-insensitive."""
cls1 = loader.resolve_pipeline_class(class_name="StableDiffusionPipeline")
cls2 = loader.resolve_pipeline_class(class_name="stablediffusionpipeline")
self.assertEqual(cls1, cls2)
def test_resolve_pipeline_class_by_task(self):
"""Test resolving pipeline class by task alias."""
# Get the registry to find available tasks
aliases = loader.get_task_aliases()
# Test with a common task that should be available
if "stable-diffusion" in aliases:
cls = loader.resolve_pipeline_class(task="stable-diffusion")
self.assertIsNotNone(cls)
def test_resolve_pipeline_class_unknown_name_raises(self):
"""Test that resolving unknown class name raises ValueError with helpful message."""
with self.assertRaises(ValueError) as ctx:
loader.resolve_pipeline_class(class_name="NonExistentPipeline")
# Check that error message includes available pipelines
error_msg = str(ctx.exception)
self.assertIn("Unknown pipeline class", error_msg)
self.assertIn("Available pipelines", error_msg)
def test_resolve_pipeline_class_unknown_task_raises(self):
"""Test that resolving unknown task raises ValueError with helpful message."""
with self.assertRaises(ValueError) as ctx:
loader.resolve_pipeline_class(task="nonexistent-task-xyz")
# Check that error message includes available tasks
error_msg = str(ctx.exception)
self.assertIn("Unknown task", error_msg)
self.assertIn("Available tasks", error_msg)
def test_resolve_pipeline_class_no_params_raises(self):
"""Test that calling with no parameters raises helpful ValueError."""
with self.assertRaises(ValueError) as ctx:
loader.resolve_pipeline_class()
error_msg = str(ctx.exception)
self.assertIn("Must provide at least one of", error_msg)
def test_get_pipeline_info(self):
"""Test getting pipeline information."""
info = loader.get_pipeline_info("StableDiffusionPipeline")
self.assertEqual(info['name'], "StableDiffusionPipeline")
self.assertIsInstance(info['aliases'], list)
self.assertIsInstance(info['supports_single_file'], bool)
def test_get_pipeline_info_unknown_raises(self):
"""Test that getting info for unknown pipeline raises ValueError."""
with self.assertRaises(ValueError) as ctx:
loader.get_pipeline_info("NonExistentPipeline")
self.assertIn("Unknown pipeline", str(ctx.exception))
def test_discover_diffusers_classes_pipelines(self):
"""Test generic class discovery for DiffusionPipeline."""
classes = loader.discover_diffusers_classes("DiffusionPipeline")
# Should return a dict
self.assertIsInstance(classes, dict)
# Should contain known pipeline classes
self.assertIn("DiffusionPipeline", classes)
self.assertIn("StableDiffusionPipeline", classes)
def test_discover_diffusers_classes_caches_results(self):
"""Test that class discovery results are cached."""
classes1 = loader.discover_diffusers_classes("DiffusionPipeline")
classes2 = loader.discover_diffusers_classes("DiffusionPipeline")
# Should be the same object (cached)
self.assertIs(classes1, classes2)
def test_discover_diffusers_classes_exclude_base(self):
"""Test discovering classes without base class."""
classes = loader.discover_diffusers_classes("DiffusionPipeline", include_base=False)
# Should still contain subclasses
self.assertIn("StableDiffusionPipeline", classes)
def test_get_available_classes(self):
"""Test getting list of available classes for a base class."""
classes = loader.get_available_classes("DiffusionPipeline")
# Should return a sorted list
self.assertIsInstance(classes, list)
self.assertEqual(classes, sorted(classes))
# Should contain known classes
self.assertIn("StableDiffusionPipeline", classes)
class TestDiffusersDynamicLoaderWithMocks(unittest.TestCase):
"""Test cases using mocks to test edge cases."""
def test_load_pipeline_requires_model_id(self):
"""Test that load_diffusers_pipeline requires model_id."""
with self.assertRaises(ValueError) as ctx:
loader.load_diffusers_pipeline(class_name="StableDiffusionPipeline")
self.assertIn("model_id is required", str(ctx.exception))
def test_resolve_with_model_id_uses_diffusion_pipeline_fallback(self):
"""Test that resolving with only model_id falls back to DiffusionPipeline."""
# When model_id is provided, if hub lookup is not successful,
# should fall back to DiffusionPipeline.
# This tests the fallback behavior - the actual hub lookup may succeed
# or fail depending on network, but the fallback path should work.
cls = loader.resolve_pipeline_class(model_id="some/nonexistent/model")
self.assertEqual(cls, DiffusionPipeline)
@unittest.skipUnless(GRPC_AVAILABLE, "gRPC modules not available")
class TestGenerateImageOptionsKwargsMerge(unittest.TestCase):
"""Test that GenerateImage merges the options dict into pipeline kwargs.
The options dict holds image (PIL), negative_prompt, and
num_inference_steps. Without the merge, img2img pipelines never
receive the source image and fail with 'Input is in incorrect format'.
"""
def test_options_merged_into_pipeline_kwargs(self):
from backend import BackendServicer
from PIL import Image
import tempfile, os
svc = BackendServicer.__new__(BackendServicer)
# Minimal attributes the method reads
svc.pipe = MagicMock()
svc.pipe.return_value.images = [Image.new("RGB", (4, 4))]
svc.cfg_scale = 7.5
svc.controlnet = None
svc.img2vid = False
svc.txt2vid = False
svc.clip_skip = 0
svc.PipelineType = "StableDiffusionImg2ImgPipeline"
svc.options = {}
# Create a tiny source image for the request's src field
src_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
Image.new("RGB", (4, 4), color="red").save(src_file, format="PNG")
src_file.close()
dst_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
dst_file.close()
try:
request = MagicMock()
request.positive_prompt = "a test prompt"
request.negative_prompt = "bad quality"
request.step = 10
request.seed = 0
request.width = 0
request.height = 0
request.src = src_file.name
request.ref_images = []
request.dst = dst_file.name
svc.GenerateImage(request, context=None)
# The pipeline must have been called with the image kwarg
svc.pipe.assert_called_once()
_, call_kwargs = svc.pipe.call_args
self.assertIn("image", call_kwargs,
"source image must be passed to pipeline via kwargs")
self.assertIn("negative_prompt", call_kwargs,
"negative_prompt must be passed to pipeline via kwargs")
self.assertEqual(call_kwargs["num_inference_steps"], 10)
finally:
os.unlink(src_file.name)
os.unlink(dst_file.name)
backend/python/diffusers/test.sh
Executable
+11
View 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
backend/python/faster-qwen3-tts/Makefile
+23
View File
@@ -0,0 +1,23 @@
.PHONY: faster-qwen3-tts
faster-qwen3-tts:
bash install.sh
.PHONY: run
run: faster-qwen3-tts
@echo "Running faster-qwen3-tts..."
bash run.sh
@echo "faster-qwen3-tts run."
.PHONY: test
test: faster-qwen3-tts
@echo "Testing faster-qwen3-tts..."
bash test.sh
@echo "faster-qwen3-tts tested."
.PHONY: protogen-clean
protogen-clean:
$(RM) backend_pb2_grpc.py backend_pb2.py
.PHONY: clean
clean: protogen-clean
rm -rf venv __pycache__
backend/python/faster-qwen3-tts/backend.py
+199
View File
@@ -0,0 +1,199 @@
#!/usr/bin/env python3
"""
gRPC server of LocalAI for Faster Qwen3-TTS (CUDA graph capture, voice clone only).
"""
from concurrent import futures
import time
import argparse
import signal
import sys
import os
import traceback
import backend_pb2
import backend_pb2_grpc
import torch
import soundfile as sf
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
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
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
MAX_WORKERS = int(os.environ.get('PYTHON_GRPC_MAX_WORKERS', '1'))
class BackendServicer(backend_pb2_grpc.BackendServicer):
def Health(self, request, context):
return backend_pb2.Reply(message=bytes("OK", 'utf-8'))
def LoadModel(self, request, context):
if not torch.cuda.is_available():
return backend_pb2.Result(
success=False,
message="faster-qwen3-tts requires NVIDIA GPU with CUDA"
)
self.options = {}
for opt in request.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"
self.options[key] = value
model_path = request.Model or "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
self.audio_path = request.AudioPath if hasattr(request, 'AudioPath') and request.AudioPath else None
self.model_file = request.ModelFile if hasattr(request, 'ModelFile') and request.ModelFile else None
self.model_path = request.ModelPath if hasattr(request, 'ModelPath') and request.ModelPath else None
from faster_qwen3_tts import FasterQwen3TTS
print(f"Loading model from: {model_path}", file=sys.stderr)
try:
self.model = FasterQwen3TTS.from_pretrained(model_path)
except Exception as e:
print(f"[ERROR] Loading model: {type(e).__name__}: {e}", file=sys.stderr)
print(traceback.format_exc(), file=sys.stderr)
return backend_pb2.Result(success=False, message=str(e))
print(f"Model loaded successfully: {model_path}", file=sys.stderr)
return backend_pb2.Result(message="Model loaded successfully", success=True)
def _get_ref_audio_path(self, request):
if not self.audio_path:
return None
if os.path.isabs(self.audio_path):
return self.audio_path
if self.model_file:
model_file_base = os.path.dirname(self.model_file)
ref_path = os.path.join(model_file_base, self.audio_path)
if os.path.exists(ref_path):
return ref_path
if self.model_path:
ref_path = os.path.join(self.model_path, self.audio_path)
if os.path.exists(ref_path):
return ref_path
return self.audio_path
def TTS(self, request, context):
try:
if not request.dst:
return backend_pb2.Result(
success=False,
message="dst (output path) is required"
)
text = request.text.strip()
if not text:
return backend_pb2.Result(
success=False,
message="Text is empty"
)
language = request.language if hasattr(request, 'language') and request.language else None
if not language or language == "":
language = "English"
ref_audio = self._get_ref_audio_path(request)
if not ref_audio:
return backend_pb2.Result(
success=False,
message="AudioPath is required for voice clone (set in LoadModel)"
)
ref_text = self.options.get("ref_text")
if not ref_text and hasattr(request, 'ref_text') and request.ref_text:
ref_text = request.ref_text
if not ref_text:
return backend_pb2.Result(
success=False,
message="ref_text is required for voice clone (set via LoadModel Options, e.g. ref_text:Your reference transcript)"
)
chunk_size = self.options.get("chunk_size")
generation_kwargs = {}
if chunk_size is not None:
generation_kwargs["chunk_size"] = int(chunk_size)
audio_list, sr = self.model.generate_voice_clone(
text=text,
language=language,
ref_audio=ref_audio,
ref_text=ref_text,
**generation_kwargs
)
if audio_list is None or (isinstance(audio_list, list) and len(audio_list) == 0):
return backend_pb2.Result(
success=False,
message="No audio output generated"
)
audio_data = audio_list[0] if isinstance(audio_list, list) else audio_list
sf.write(request.dst, audio_data, sr)
print(f"Saved output to {request.dst}", file=sys.stderr)
except Exception as err:
print(f"Error in TTS: {err}", file=sys.stderr)
print(traceback.format_exc(), file=sys.stderr)
return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")
return backend_pb2.Result(success=True)
def serve(address):
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("Server started. Listening on: " + address, file=sys.stderr)
def signal_handler(sig, frame):
print("Received termination signal. Shutting down...")
server.stop(0)
sys.exit(0)
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
try:
while True:
time.sleep(_ONE_DAY_IN_SECONDS)
except KeyboardInterrupt:
server.stop(0)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Run the gRPC server.")
parser.add_argument("--addr", default="localhost:50051", help="The address to bind the server to.")
args = parser.parse_args()
serve(args.addr)
backend/python/faster-qwen3-tts/install.sh
+13
View File
@@ -0,0 +1,13 @@
#!/bin/bash
set -e
EXTRA_PIP_INSTALL_FLAGS="--no-build-isolation"
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
backend/python/faster-qwen3-tts/requirements-cublas12.txt
+4
View File
@@ -0,0 +1,4 @@
--extra-index-url https://download.pytorch.org/whl/cu121
torch
torchaudio
faster-qwen3-tts
backend/python/faster-qwen3-tts/requirements-cublas13.txt
+4
View File
@@ -0,0 +1,4 @@
--extra-index-url https://download.pytorch.org/whl/cu130
torch
torchaudio
faster-qwen3-tts
backend/python/faster-qwen3-tts/requirements-install.txt
+1
View File
@@ -0,0 +1 @@
setuptools
backend/python/faster-qwen3-tts/requirements-l4t12.txt
+4
View File
@@ -0,0 +1,4 @@
--extra-index-url https://pypi.jetson-ai-lab.io/jp6/cu129/
torch
torchaudio
faster-qwen3-tts
backend/python/faster-qwen3-tts/requirements-l4t13.txt
+4
View File
@@ -0,0 +1,4 @@
--extra-index-url https://download.pytorch.org/whl/cu130
torch
torchaudio
faster-qwen3-tts
backend/python/faster-qwen3-tts/requirements.txt
+9
View File
@@ -0,0 +1,9 @@
grpcio==1.71.0
protobuf
certifi
packaging==24.1
soundfile
setuptools
six
anyio
sox
backend/python/faster-qwen3-tts/run.sh
+9
View 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 $@
backend/python/faster-qwen3-tts/test.py
+104
View File
@@ -0,0 +1,104 @@
"""
Tests for the faster-qwen3-tts gRPC backend.
"""
import unittest
import subprocess
import time
import os
import sys
import tempfile
import backend_pb2
import backend_pb2_grpc
import grpc
class TestBackendServicer(unittest.TestCase):
def setUp(self):
self.service = subprocess.Popen(
["python3", "backend.py", "--addr", "localhost:50052"],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
cwd=os.path.dirname(os.path.abspath(__file__)),
)
time.sleep(15)
def tearDown(self):
self.service.terminate()
try:
self.service.communicate(timeout=5)
except subprocess.TimeoutExpired:
self.service.kill()
self.service.communicate()
def test_health(self):
with grpc.insecure_channel("localhost:50052") as channel:
stub = backend_pb2_grpc.BackendStub(channel)
reply = stub.Health(backend_pb2.HealthMessage(), timeout=5.0)
self.assertEqual(reply.message, b"OK")
def test_load_model_requires_cuda(self):
with grpc.insecure_channel("localhost:50052") as channel:
stub = backend_pb2_grpc.BackendStub(channel)
response = stub.LoadModel(
backend_pb2.ModelOptions(
Model="Qwen/Qwen3-TTS-12Hz-0.6B-Base",
CUDA=True,
),
timeout=10.0,
)
self.assertFalse(response.success)
@unittest.skipUnless(
__import__("torch").cuda.is_available(),
"faster-qwen3-tts TTS requires CUDA",
)
def test_tts(self):
import soundfile as sf
try:
with grpc.insecure_channel("localhost:50052") as channel:
stub = backend_pb2_grpc.BackendStub(channel)
ref_audio = tempfile.NamedTemporaryFile(suffix='.wav', delete=False)
ref_audio.close()
try:
sr = 22050
duration = 1.0
samples = int(sr * duration)
sf.write(ref_audio.name, [0.0] * samples, sr)
response = stub.LoadModel(
backend_pb2.ModelOptions(
Model="Qwen/Qwen3-TTS-12Hz-0.6B-Base",
AudioPath=ref_audio.name,
Options=["ref_text:Hello world"],
),
timeout=600.0,
)
self.assertTrue(response.success, response.message)
with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as out:
output_path = out.name
try:
tts_response = stub.TTS(
backend_pb2.TTSRequest(
text="Test output.",
dst=output_path,
language="English",
),
timeout=120.0,
)
self.assertTrue(tts_response.success, tts_response.message)
self.assertTrue(os.path.exists(output_path))
self.assertGreater(os.path.getsize(output_path), 0)
finally:
if os.path.exists(output_path):
os.unlink(output_path)
finally:
if os.path.exists(ref_audio.name):
os.unlink(ref_audio.name)
except Exception as err:
self.fail(f"TTS test failed: {err}")
if __name__ == "__main__":
unittest.main()
backend/python/faster-qwen3-tts/test.sh
+11
View 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
backend/python/faster-whisper/Makefile
+14
View File
@@ -0,0 +1,14 @@
.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__
# trigger per-arch+merge rebuild for faster-whisper pilot
backend/python/faster-whisper/backend.py
Executable
+124
View File
@@ -0,0 +1,124 @@
#!/usr/bin/env python3
"""
This is an extra gRPC server of LocalAI for Faster Whisper TTS
"""
from concurrent import futures
import time
import argparse
import signal
import sys
import os
import backend_pb2
import backend_pb2_grpc
import torch
from faster_whisper import WhisperModel
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
_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'))
COQUI_LANGUAGE = os.environ.get('COQUI_LANGUAGE', None)
# Implement the BackendServicer class with the service methods
class BackendServicer(backend_pb2_grpc.BackendServicer):
"""
BackendServicer is the class that implements the gRPC service
"""
def Health(self, request, context):
return backend_pb2.Reply(message=bytes("OK", 'utf-8'))
def LoadModel(self, request, context):
device = "cpu"
# Get device
# device = "cuda" if request.CUDA else "cpu"
if request.CUDA:
device = "cuda"
mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
if mps_available:
device = "mps"
try:
print("Preparing models, please wait", file=sys.stderr)
self.model = WhisperModel(request.Model, device=device, compute_type="default")
except Exception as err:
return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")
# Implement your logic here for the LoadModel service
# Replace this with your desired response
return backend_pb2.Result(message="Model loaded successfully", success=True)
def AudioTranscription(self, request, context):
resultSegments = []
text = ""
try:
word_timestamps = "word" in request.timestamp_granularities
segments, info = self.model.transcribe(request.dst, beam_size=5, condition_on_previous_text=False, word_timestamps=word_timestamps)
id = 0
for segment in segments:
print("[%.2fs -> %.2fs] %s" % (segment.start, segment.end, segment.text))
words = []
if word_timestamps and hasattr(segment, 'words'):
for word in segment.words:
words.append(backend_pb2.TranscriptWord(
start=int(word.start * 1e9),
end=int(word.end * 1e9),
text=word.word
))
resultSegments.append(backend_pb2.TranscriptSegment(
id=id,
start=int(segment.start * 1e9),
end=int(segment.end * 1e9),
text=segment.text,
words=words
))
text += segment.text
id += 1
except Exception as err:
print(f"Unexpected {err=}, {type(err)=}", file=sys.stderr)
raise err
return backend_pb2.TranscriptResult(segments=resultSegments, text=text)
def serve(address):
server = grpc.server(futures.ThreadPoolExecutor(max_workers=MAX_WORKERS),
options=[
('grpc.max_message_length', 50 * 1024 * 1024), # 50MB
('grpc.max_send_message_length', 50 * 1024 * 1024), # 50MB
('grpc.max_receive_message_length', 50 * 1024 * 1024), # 50MB
],
interceptors=get_auth_interceptors(),
)
backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
server.add_insecure_port(address)
server.start()
print("Server started. Listening on: " + address, file=sys.stderr)
# Define the signal handler function
def signal_handler(sig, frame):
print("Received termination signal. Shutting down...")
server.stop(0)
sys.exit(0)
# Set the signal handlers for SIGINT and SIGTERM
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
try:
while True:
time.sleep(_ONE_DAY_IN_SECONDS)
except KeyboardInterrupt:
server.stop(0)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Run the gRPC server.")
parser.add_argument(
"--addr", default="localhost:50051", help="The address to bind the server to."
)
args = parser.parse_args()
serve(args.addr)
backend/python/faster-whisper/install.sh
Executable
+29
View File
@@ -0,0 +1,29 @@
#!/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
# This is here because the Intel pip index is broken and returns 200 status codes for every package name, it just doesn't return any package links.
# This makes uv think that the package exists in the Intel pip index, and by default it stops looking at other pip indexes once it finds a match.
# We need uv to continue falling through to the pypi default index to find optimum[openvino] in the pypi index
# the --upgrade actually allows us to *downgrade* torch to the version provided in the Intel pip index
if [ "x${BUILD_PROFILE}" == "xintel" ]; then
EXTRA_PIP_INSTALL_FLAGS+=" --upgrade --index-strategy=unsafe-first-match"
fi
if [ "x${BUILD_PROFILE}" == "xl4t13" ]; then
PYTHON_VERSION="3.12"
PYTHON_PATCH="12"
PY_STANDALONE_TAG="20251120"
fi
if [ "x${BUILD_PROFILE}" == "xl4t12" ]; then
USE_PIP=true
fi
installRequirements
backend/python/faster-whisper/protogen.sh
+11
View 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
python3 -m grpc_tools.protoc -I../.. -I./ --python_out=. --grpc_python_out=. backend.proto
backend/python/faster-whisper/requirements-cpu.txt
+8
View File
@@ -0,0 +1,8 @@
faster-whisper
opencv-python
accelerate
compel
peft
sentencepiece
torch==2.4.1
optimum-quanto
backend/python/faster-whisper/requirements-cublas12.txt
+8
View File
@@ -0,0 +1,8 @@
torch==2.4.1
faster-whisper
opencv-python
accelerate
compel
peft
sentencepiece
optimum-quanto
backend/python/faster-whisper/requirements-cublas13.txt
+9
View File
@@ -0,0 +1,9 @@
--extra-index-url https://download.pytorch.org/whl/cu130
torch==2.9.1
faster-whisper
opencv-python
accelerate
compel
peft
sentencepiece
optimum-quanto
backend/python/faster-whisper/requirements-hipblas.txt
+3
View File
@@ -0,0 +1,3 @@
--extra-index-url https://download.pytorch.org/whl/rocm7.0
torch
faster-whisper
backend/python/faster-whisper/requirements-intel.txt
+4
View File
@@ -0,0 +1,4 @@
--extra-index-url https://download.pytorch.org/whl/xpu
torch
optimum[openvino]
faster-whisper

Some files were not shown because too many files have changed in this diff Show More