961 lines · 31.8 KB · TypeScript
RawDownload
import { constants, existsSync, statfsSync } from "node:fs";import { open, unlink } from "node:fs/promises";import { join, relative, resolve } from "node:path";import { Effect, Semaphore } from "effect";import {  CHATTERBOX_BACKEND,  CHATTERBOX_MODEL_REVISION,  CHATTERBOX_PACKAGE_VERSION,  type SpeechGpuTarget,  type SpeechStatus,  type SpeechVoiceProfile,} from "@local-studio/contracts/speech";import type { ProcessInfo, Recipe, GpuInfo } from "../models/types";import {  GpuLeaseConflict,  type GpuLeaseRegistry,  resolveRecipeGpuUuids,} from "../system/gpu-leases";import { resolveBinary } from "../../core/command";import {  ChatterboxRuntime,  chatterboxRuntimePaths,  type ChatterboxInstallOptions,  type ChatterboxRuntimeState,} from "./runtime";import {  ChatterboxWorkerClient,  type ChatterboxSynthesisInput,  type ChatterboxSynthesisResult,} from "./worker-client";import {  normalizeVoiceReference,  VoiceReferenceError,  type NormalizedVoiceReference,} from "./reference-audio";import { VoiceStore, type VoiceProfile } from "./voice-store";import { secureSpeechDirectory } from "./storage";import { queryNvidiaComputeGpuUuids } from "../system/platform/nvidia-compute-processes"; const FULL_NVIDIA_UUID =  /^GPU-[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$/;const RTX_3090_NAME = /\bRTX\s+3090\b/i;const MANAGED_INSTALL_BYTES = 32 * 1024 ** 3;const MINIMUM_FREE_RESERVE_BYTES = 8 * 1024 ** 3;const REQUIRED_INSTALL_BYTES = MANAGED_INSTALL_BYTES + MINIMUM_FREE_RESERVE_BYTES;const MAXIMUM_OUTPUT_BYTES = 32 * 1024 * 1024;const MAXIMUM_QUEUED_SYNTHESIS = 4;const MAXIMUM_PENDING_NORMALIZATION = 2;const MAXIMUM_TEXT_CHARACTERS = 4096; export class SpeechServiceError extends Error {  constructor(    readonly status: number,    readonly code: string,    message: string,  ) {    super(message);    this.name = "SpeechServiceError";  }} export interface SpeechEngineState {  getCurrentProcess(): Promise<ProcessInfo | null>;  getCurrentRecipe(): Promise<Recipe | null>;} export interface SpeechRuntime {  readonly paths: { readonly pythonPath: string };  getState(): ChatterboxRuntimeState;  startInstall(gpuUuid: string, options?: ChatterboxInstallOptions): ChatterboxRuntimeState;  waitForInstall(): Promise<ChatterboxRuntimeState>;  cancelInstall(): Promise<void>;} export interface SpeechWorker {  synthesize(input: ChatterboxSynthesisInput): Promise<ChatterboxSynthesisResult>;  shutdown(): Promise<void>;  settleTermination(): Promise<void>;  terminate(): Promise<void>;} const speechGpuLeaseBrand: unique symbol = Symbol("SpeechGpuLeaseGuard"); export interface SpeechGpuLeaseGuard {  readonly uuid: string;  readonly generation: number;  readonly [speechGpuLeaseBrand]: true;} export interface SpeechVoiceStore {  list(): VoiceProfile[];  create(input: {    name: string;    durationMs: number;    consent: string;    audio: Uint8Array;  }): Promise<VoiceProfile>;  delete(id: string): Promise<boolean>;  withPlaintext<A>(id: string, use: (path: string) => Promise<A>): Promise<A>;} export interface SpeechDiskAvailability {  readonly totalBytes: number;  readonly availableBytes: number;} export interface SpeechSynthesisInput {  readonly text: string;  readonly voiceId: string;} export interface SpeechSynthesisOutput {  readonly audio: Uint8Array;  readonly contentType: "audio/wav";  readonly sampleRate: number;} export interface SpeechVoiceInput {  readonly name: string;  readonly consent: string;  readonly audio: Uint8Array;} export interface SpeechInstallInput {  readonly repair?: boolean | undefined;} export interface SpeechServiceOptions {  readonly dataDirectory: string;  readonly databasePath: string;  readonly engine: SpeechEngineState;  readonly gpuLeaseRegistry: GpuLeaseRegistry;  readonly gpuInfo: () => GpuInfo[];  readonly environment?: NodeJS.ProcessEnv | undefined;  readonly runtime?: SpeechRuntime | undefined;  readonly voiceStore?: SpeechVoiceStore | undefined;  readonly workerFactory?: ((lease: SpeechGpuLeaseGuard) => SpeechWorker) | undefined;  readonly normalizeReference?:    ((input: Uint8Array, dataDirectory: string) => Promise<NormalizedVoiceReference>) | undefined;  readonly diskAvailability?: (() => SpeechDiskAvailability | null) | undefined;  readonly resolveBinary?: ((name: string) => string | null) | undefined;  readonly computeGpuUuids?: (() => Promise<readonly string[]>) | undefined;} const canonicalUuid = (uuid: string): string => `GPU-${uuid.slice(4).toLowerCase()}`; const serviceError = (error: unknown, status = 500, code = "speech_failed"): SpeechServiceError =>  error instanceof SpeechServiceError    ? error    : new SpeechServiceError(status, code, error instanceof Error ? error.message : String(error)); const installationMessage = (state: ChatterboxRuntimeState): string => {  if (state.status === "not_installed") return "Chatterbox Turbo is not installed";  if (state.status === "installed") return "Chatterbox Turbo is ready";  if (state.status === "error") return state.message;  if (state.stage === "preparing") return "Preparing Chatterbox Turbo";  if (state.stage === "creating_runtime") return "Creating the speech runtime";  if (state.stage === "installing_package") return "Installing Chatterbox Turbo";  return "Downloading the pinned Chatterbox Turbo model";}; const installationStatus = (state: ChatterboxRuntimeState): SpeechStatus["install"] => {  if (state.status === "not_installed") {    return { phase: "missing", progress: 0, message: installationMessage(state), error: null };  }  if (state.status === "installed") {    return { phase: "ready", progress: 1, message: installationMessage(state), error: null };  }  if (state.status === "error") {    return { phase: "failed", progress: 0, message: state.message, error: state.message };  }  return {    phase: "installing",    progress: state.progress,    message: installationMessage(state),    error: null,  };}; const diskAvailability = (path: string): SpeechDiskAvailability | null => {  try {    const stats = statfsSync(path);    return {      totalBytes: stats.blocks * stats.bsize,      availableBytes: stats.bavail * stats.bsize,    };  } catch {    return null;  }}; const outputChildPath = (directory: string, path: string): string => {  const root = resolve(directory);  const candidate = resolve(path);  const child = relative(root, candidate);  if (!child || child.startsWith("..") || child.startsWith("/") || child.startsWith("\\")) {    throw new SpeechServiceError(502, "speech_output_invalid", "Speech worker output is invalid");  }  return candidate;}; const validatedWave = (audio: Uint8Array): Uint8Array => {  const bytes = Buffer.from(audio);  if (    bytes.length < 44 ||    bytes.subarray(0, 4).toString("ascii") !== "RIFF" ||    bytes.subarray(8, 12).toString("ascii") !== "WAVE" ||    bytes.readUInt32LE(4) + 8 !== bytes.length  ) {    throw new SpeechServiceError(502, "speech_output_invalid", "Speech worker output is invalid");  }  return bytes;}; const readBoundedWave = async (path: string): Promise<Uint8Array> => {  const handle = await open(path, constants.O_RDONLY | constants.O_NOFOLLOW);  try {    const stats = await handle.stat();    if (!stats.isFile() || stats.size > MAXIMUM_OUTPUT_BYTES) {      throw new SpeechServiceError(502, "speech_output_invalid", "Speech worker output is invalid");    }    const bytes = Buffer.alloc(Math.min(MAXIMUM_OUTPUT_BYTES + 1, stats.size + 1));    let offset = 0;    while (offset < bytes.length) {      const result = await handle.read(bytes, offset, bytes.length - offset, offset);      if (result.bytesRead === 0) break;      offset += result.bytesRead;    }    if (offset > MAXIMUM_OUTPUT_BYTES) {      throw new SpeechServiceError(502, "speech_output_invalid", "Speech worker output is invalid");    }    const completed = await handle.stat();    if (completed.size !== offset) {      throw new SpeechServiceError(502, "speech_output_invalid", "Speech worker output is invalid");    }    return validatedWave(Buffer.from(bytes.subarray(0, offset)));  } finally {    await handle.close();  }}; const validText = (text: string): string => {  if (!text.trim())    throw new SpeechServiceError(400, "speech_text_required", "Speech text is required");  if (Array.from(text).length > MAXIMUM_TEXT_CHARACTERS) {    throw new SpeechServiceError(      400,      "speech_text_too_long",      `Speech text cannot exceed ${MAXIMUM_TEXT_CHARACTERS} characters`,    );  }  return text;}; const stoppingError = (): SpeechServiceError =>  new SpeechServiceError(409, "speech_stopping", "Speech runtime is stopping"); export class SpeechService {  private readonly dataDirectory: string;  private readonly environment: NodeJS.ProcessEnv;  private readonly runtime: SpeechRuntime;  private readonly voiceStore: SpeechVoiceStore;  private readonly workerFactory: (lease: SpeechGpuLeaseGuard) => SpeechWorker;  private readonly normalizeReference: (    input: Uint8Array,    dataDirectory: string,  ) => Promise<NormalizedVoiceReference>;  private readonly getDiskAvailability: () => SpeechDiskAvailability | null;  private readonly findBinary: (name: string) => string | null;  private readonly computeGpuUuids: () => Promise<readonly string[]>;  private readonly outputDirectory: string;  private readonly activation = Semaphore.makeUnsafe(1);  private readonly synthesis = Semaphore.makeUnsafe(1);  private readonly voiceNormalization = Semaphore.makeUnsafe(1);  private worker: SpeechWorker | null = null;  private workerPhase: SpeechStatus["worker"]["phase"] = "stopped";  private workerError: string | null = null;  private leasedGpuUuid: string | null = null;  private liveLease: SpeechGpuLeaseGuard | null = null;  private leaseGeneration = 0;  private quarantined = false;  private pendingSynthesis = 0;  private pendingNormalization = 0;  private acceptingSynthesis = true;  private synthesisEpoch = 0;  private installTask: Promise<void> | null = null;  private cancellingInstall = false;  private stopTask: Promise<void> | null = null;   constructor(private readonly options: SpeechServiceOptions) {    this.environment = options.environment ?? process.env;    this.dataDirectory = resolve(      this.environment["LOCAL_STUDIO_SPEECH_DATA_DIR"] ?? options.dataDirectory,    );    const paths = chatterboxRuntimePaths(this.dataDirectory);    this.outputDirectory = paths.outputDirectory;    secureSpeechDirectory(paths.speechDirectory);    this.runtime = options.runtime ?? new ChatterboxRuntime({ dataDirectory: this.dataDirectory });    this.voiceStore =      options.voiceStore ?? new VoiceStore(options.databasePath, this.dataDirectory);    this.workerFactory =      options.workerFactory ??      ((lease): SpeechWorker =>        new ChatterboxWorkerClient({          dataDirectory: this.dataDirectory,          gpuUuid: lease.uuid,          voiceDirectory: join(this.dataDirectory, "runtime", "speech", "tmp"),        }));    this.normalizeReference = options.normalizeReference ?? normalizeVoiceReference;    this.getDiskAvailability =      options.diskAvailability ??      ((): SpeechDiskAvailability | null => diskAvailability(this.dataDirectory));    this.findBinary = options.resolveBinary ?? resolveBinary;    this.computeGpuUuids = options.computeGpuUuids ?? queryNvidiaComputeGpuUuids;  }   getStatus(): SpeechStatus {    const target = this.statusTarget();    const storage = this.getDiskAvailability();    return {      backend: CHATTERBOX_BACKEND,      package_version: CHATTERBOX_PACKAGE_VERSION,      model_revision: CHATTERBOX_MODEL_REVISION,      install: installationStatus(this.runtime.getState()),      worker: {        phase: this.workerPhase,        queue_depth: Math.max(0, this.pendingSynthesis - 1),        error: this.workerError,      },      gpu: target,      prerequisites: {        ffmpeg: Boolean(this.findBinary(this.environment["LOCAL_STUDIO_FFMPEG_CLI"] ?? "ffmpeg")),        python_311: Boolean(          existsSync(this.runtime.paths.pythonPath) ||          this.findBinary("python3.11") ||          this.findBinary("uv"),        ),        storage: {          available_bytes: storage?.availableBytes ?? null,          required_bytes: REQUIRED_INSTALL_BYTES,          ready: Boolean(storage && storage.availableBytes >= REQUIRED_INSTALL_BYTES),        },      },      voice_count: this.voiceStore.list().length,    };  }   install(input: SpeechInstallInput = {}): Promise<SpeechStatus> {    if (this.stopTask) return Promise.reject(stoppingError());    const state = this.runtime.getState();    if (state.status === "installing" || (state.status === "installed" && !input.repair)) {      return Promise.resolve(this.getStatus());    }    return Effect.runPromise(      this.activation.withPermit(        Effect.tryPromise({          try: async () => {            const current = this.runtime.getState();            if (              current.status === "installing" ||              (current.status === "installed" && !input.repair)            ) {              return this.getStatus();            }            if (input.repair && this.worker) {              await this.stopRuntime(false, true);            }            this.assertInstallCapacity();            const lease = await this.activateSpeech();            await this.assertLiveLease(lease);            let started: ChatterboxRuntimeState;            try {              started = this.startRuntimeInstall(lease, input);            } catch (error) {              if (!this.worker) await this.releaseSpeechLease();              throw error;            }            if (started.status !== "installing") {              if (!this.worker) await this.releaseSpeechLease();              if (started.status === "error") {                throw new SpeechServiceError(500, "speech_install_failed", started.message);              }              return this.getStatus();            }            this.startInstallCompletion(lease);            return this.getStatus();          },          catch: (error) => serviceError(error, 500, "speech_install_failed"),        }),      ),    );  }   listVoices(): SpeechVoiceProfile[] {    return this.voiceStore.list();  }   createVoice(input: SpeechVoiceInput): Promise<SpeechVoiceProfile> {    if (this.pendingNormalization >= MAXIMUM_PENDING_NORMALIZATION) {      return Promise.reject(        new VoiceReferenceError(429, "voice_queue_full", "Voice normalization queue is full"),      );    }    this.pendingNormalization += 1;    return Effect.runPromise(      this.voiceNormalization        .withPermit(          Effect.tryPromise({            try: async () => {              const normalized = await this.normalizeReference(input.audio, this.dataDirectory);              return this.voiceStore.create({                name: input.name,                consent: input.consent,                audio: normalized.audio,                durationMs: normalized.durationMs,              });            },            catch: (error) => error,          }),        )        .pipe(          Effect.ensuring(            Effect.sync(() => {              this.pendingNormalization -= 1;            }),          ),        ),    );  }   deleteVoice(id: string): Promise<boolean> {    return this.voiceStore.delete(id);  }   synthesize(input: SpeechSynthesisInput): Promise<SpeechSynthesisOutput> {    if (!this.acceptingSynthesis) {      return Promise.reject(stoppingError());    }    if (this.pendingSynthesis >= MAXIMUM_QUEUED_SYNTHESIS + 1) {      return Promise.reject(        new SpeechServiceError(429, "speech_queue_full", "Speech queue is full"),      );    }    this.pendingSynthesis += 1;    const epoch = this.synthesisEpoch;    const operation = this.synthesis.withPermit(      Effect.suspend(() =>        epoch === this.synthesisEpoch          ? Effect.tryPromise({              try: () => this.synthesizeOne(input, epoch),              catch: (error) => error,            })          : Effect.fail(stoppingError()),      ),    );    return Effect.runPromise(      operation.pipe(        Effect.ensuring(          Effect.sync(() => {            this.pendingSynthesis -= 1;          }),        ),      ),    );  }   stop(): Promise<void> {    if (this.installTask) {      return Promise.reject(        new SpeechServiceError(          409,          "speech_installing",          "Wait for the Chatterbox install to finish before stopping speech",        ),      );    }    return this.stopRuntime(false, true);  }   cancelInstall(): Promise<void> {    if (!this.installTask) return Promise.resolve();    return this.stopRuntime(true, true);  }   shutdown(): Promise<void> {    return this.stopRuntime(true, false);  }   private statusTarget(): SpeechGpuTarget | null {    try {      return this.resolveTarget(this.options.gpuInfo());    } catch {      return null;    }  }   private resolveTarget(gpus: readonly GpuInfo[]): SpeechGpuTarget {    if (gpus.length === 0) {      throw new SpeechServiceError(        503,

Lines 1–500 of 961.