mirror of
https://github.com/JezzWTF/vibepod.git
synced 2026-06-01 15:22:14 +00:00
LyAhn
7591d15a52
Overlap acoustic_decode with forward_tts_lm calls using a background ThreadPoolExecutor, hiding ~72s of decode cost behind tts_lm work. Achieved 0.67x realtime (up from 0.43x, ~56% improvement). - vibevoice_generate_patch.py: patched generate() loop reordered to submit decode to thread before running connector + tts_lm×2, then resolve future. Installed as instance method via types.MethodType so uv sync reinstalling the package cannot revert the patch. - Dynamic INT8 quantization of Linear layers (VIBEPOD_QUANTIZE=1, default on CPU). prediction_head excluded — small fixed-size tensors regressed ~20% with INT8 due to pack/unpack overhead. - Auto-detect AVX512_BF16 and load model in bfloat16 if supported (VIBEPOD_CPU_BF16=auto, overridable with 0/1). - CPU thread count auto-configured from logical CPU count; OMP/MKL env vars set accordingly. Lock file preserved around uv sync --no-sources so CPU mode does not alter the shared uv.lock. - torch.compile retained as opt-in (VIBEPOD_COMPILE=1) but marked not recommended — dynamic KV cache shapes prevent kernel reuse.
164 lines
6.5 KiB
Bash
Executable File
164 lines
6.5 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
# VibePod TTS server — start script
|
|
# Syncs the uv environment, downloads the model on first run, then launches uvicorn.
|
|
# Prerequisite: uv must be installed (https://docs.astral.sh/uv/getting-started/installation/)
|
|
#
|
|
# Usage:
|
|
# ./start.sh — CUDA mode (default, uses PyTorch CUDA 12.4 wheel, venv: .venv)
|
|
# ./start.sh --cpu — CPU-only mode (uses PyPI CPU torch wheel, venv: .venv-cpu)
|
|
#
|
|
# Optional CUDA acceleration:
|
|
# VIBEPOD_ENABLE_FLASH_ATTN=1 ./start.sh
|
|
# Installs a matching third-party Windows flash-attn wheel when the CUDA venv
|
|
# uses Python 3.12, torch 2.6.0, and CUDA 12.4.
|
|
#
|
|
# The two modes maintain completely separate virtual environments so their torch
|
|
# installations never conflict. UV_PROJECT_ENVIRONMENT tells uv which venv to use;
|
|
# --no-sources skips [tool.uv.sources] so the CPU run pulls the default PyPI torch wheel.
|
|
|
|
set -euo pipefail
|
|
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
|
cd "$SCRIPT_DIR"
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Parse flags
|
|
# ---------------------------------------------------------------------------
|
|
CPU_MODE=false
|
|
PASSTHROUGH_ARGS=()
|
|
|
|
for arg in "$@"; do
|
|
case "$arg" in
|
|
--cpu) CPU_MODE=true ;;
|
|
*) PASSTHROUGH_ARGS+=("$arg") ;;
|
|
esac
|
|
done
|
|
|
|
echo "================================================"
|
|
echo " VibePod TTS Server"
|
|
if $CPU_MODE; then
|
|
echo " Mode : CPU-only"
|
|
else
|
|
echo " Mode : CUDA (default)"
|
|
fi
|
|
echo "================================================"
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 1. Check uv is available
|
|
# ---------------------------------------------------------------------------
|
|
if ! command -v uv &>/dev/null; then
|
|
echo ""
|
|
echo "ERROR: uv is not installed."
|
|
echo "Install it first:"
|
|
echo " Windows: winget install astral-sh.uv"
|
|
echo " macOS/Linux: curl -LsSf https://astral.sh/uv/install.sh | sh"
|
|
echo ""
|
|
exit 1
|
|
fi
|
|
|
|
validate_flash_attn() {
|
|
uv run python -c "import flash_attn; import triton; import transformers.modeling_utils" &>/dev/null
|
|
}
|
|
|
|
remove_broken_flash_attn() {
|
|
if uv run python -c "import importlib.util; raise SystemExit(0 if importlib.util.find_spec('flash_attn') else 1)" &>/dev/null; then
|
|
if ! validate_flash_attn; then
|
|
echo " Installed flash-attn is not usable in this environment; removing it."
|
|
uv pip uninstall flash-attn
|
|
fi
|
|
fi
|
|
}
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 2. Sync Python environment
|
|
# CPU mode: use .venv-cpu and skip [tool.uv.sources] so uv pulls the
|
|
# default PyPI (CPU-only) torch wheel instead of the CUDA one.
|
|
# CUDA mode: standard uv sync — uses .venv and respects [tool.uv.sources].
|
|
# ---------------------------------------------------------------------------
|
|
echo ""
|
|
if $CPU_MODE; then
|
|
echo "--> Syncing CPU Python environment (.venv-cpu)..."
|
|
export UV_PROJECT_ENVIRONMENT=".venv-cpu"
|
|
LOCK_BACKUP=""
|
|
if [[ -f uv.lock ]]; then
|
|
LOCK_BACKUP="$(mktemp)"
|
|
cp uv.lock "$LOCK_BACKUP"
|
|
fi
|
|
uv sync --no-sources
|
|
if [[ -n "$LOCK_BACKUP" ]]; then
|
|
cp "$LOCK_BACKUP" uv.lock
|
|
rm -f "$LOCK_BACKUP"
|
|
fi
|
|
else
|
|
echo "--> Syncing CUDA Python environment (.venv)..."
|
|
uv sync
|
|
|
|
remove_broken_flash_attn
|
|
|
|
if [[ "${VIBEPOD_ENABLE_FLASH_ATTN:-0}" == "1" ]]; then
|
|
echo ""
|
|
echo "--> Checking optional FlashAttention wheel..."
|
|
|
|
if validate_flash_attn; then
|
|
echo " flash-attn already installed and importable."
|
|
else
|
|
PY_TAG="$(uv run python -c "import sys; print(f'cp{sys.version_info.major}{sys.version_info.minor}')")"
|
|
TORCH_TAG="$(uv run python -c "import torch; print(torch.__version__.split('+', 1)[0])")"
|
|
CUDA_TAG="$(uv run python -c "import torch; print('cu' + torch.version.cuda.replace('.', ''))")"
|
|
|
|
if [[ "$PY_TAG" == "cp312" && "$TORCH_TAG" == "2.6.0" && "$CUDA_TAG" == "cu124" ]]; then
|
|
FLASH_ATTN_WHEEL_URL="https://huggingface.co/lldacing/flash-attention-windows-wheel/resolve/main/flash_attn-2.7.4%2Bcu124torch2.6.0cxx11abiFALSE-cp312-cp312-win_amd64.whl"
|
|
echo " Installing flash-attn for Python 3.12, torch 2.6.0, CUDA 12.4..."
|
|
uv pip install "$FLASH_ATTN_WHEEL_URL"
|
|
if validate_flash_attn; then
|
|
echo " flash-attn import check passed."
|
|
else
|
|
echo " flash-attn import check failed; removing it and continuing with SDPA."
|
|
uv pip uninstall flash-attn
|
|
fi
|
|
else
|
|
echo " No known wheel for Python tag $PY_TAG, torch $TORCH_TAG, CUDA $CUDA_TAG."
|
|
echo " Continuing with PyTorch SDPA attention."
|
|
fi
|
|
fi
|
|
fi
|
|
fi
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 3. Launch uvicorn
|
|
# Pass DEVICE env var so the server can select the correct torch device.
|
|
# ---------------------------------------------------------------------------
|
|
echo ""
|
|
echo "--> Starting uvicorn on http://127.0.0.1:8000"
|
|
export PYTHONUTF8=1
|
|
|
|
if $CPU_MODE; then
|
|
export VIBEPOD_DEVICE="cpu"
|
|
export UV_PROJECT_ENVIRONMENT=".venv-cpu"
|
|
if [[ -z "${VIBEPOD_CPU_THREADS:-}" ]]; then
|
|
VIBEPOD_CPU_THREADS="$(uv run --no-sources python -c "import os; print(max(1, (os.cpu_count() or 2) // 2))")"
|
|
export VIBEPOD_CPU_THREADS
|
|
fi
|
|
export OMP_NUM_THREADS="${OMP_NUM_THREADS:-$VIBEPOD_CPU_THREADS}"
|
|
export MKL_NUM_THREADS="${MKL_NUM_THREADS:-$VIBEPOD_CPU_THREADS}"
|
|
# Dynamic INT8 quantization — on by default for CPU (~22% faster, prediction_head
|
|
# excluded automatically to avoid regression on small fixed-size tensors).
|
|
# Set VIBEPOD_QUANTIZE=0 to disable if you notice audio quality differences.
|
|
export VIBEPOD_QUANTIZE="${VIBEPOD_QUANTIZE:-1}"
|
|
# Optional CPU flags:
|
|
# VIBEPOD_ASYNC_DECODE=0 Disable async decode/tts_lm overlap (on by default)
|
|
# VIBEPOD_CPU_BF16=1 Force bfloat16 weights (auto-detected via AVX512_BF16)
|
|
# VIBEPOD_COMPILE=1 torch.compile hot paths (ineffective for autoregressive
|
|
# models on CPU — not recommended, kept for experimentation)
|
|
UV_RUN_ARGS=(--no-sync --no-sources)
|
|
else
|
|
export VIBEPOD_DEVICE="cuda"
|
|
UV_RUN_ARGS=()
|
|
fi
|
|
|
|
exec uv run "${UV_RUN_ARGS[@]}" uvicorn vibevoice_server:app \
|
|
--host 127.0.0.1 \
|
|
--port 8000 \
|
|
--log-level info \
|
|
"${PASSTHROUGH_ARGS[@]+"${PASSTHROUGH_ARGS[@]}"}"
|