mirror of
https://github.com/mudler/LocalAI.git
synced 2024-06-07 19:40:48 +00:00
ee59e7d45f
**Description** This PR fixes # **Notes for Reviewers** **[Signed commits](../CONTRIBUTING.md#signing-off-on-commits-developer-certificate-of-origin)** - [ ] Yes, I signed my commits. <!-- Thank you for contributing to LocalAI! Contributing Conventions: 1. Include descriptive PR titles with [<component-name>] prepended. 2. Build and test your changes before submitting a PR. 3. Sign your commits By following the community's contribution conventions upfront, the review process will be accelerated and your PR merged more quickly. --> Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
97 lines
3.8 KiB
Python
97 lines
3.8 KiB
Python
#!/usr/bin/env python3
|
|
import grpc
|
|
from concurrent import futures
|
|
import time
|
|
import backend_pb2
|
|
import backend_pb2_grpc
|
|
import argparse
|
|
import signal
|
|
import sys
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from utils.generation import SAMPLE_RATE, generate_audio, preload_models
|
|
from scipy.io.wavfile import write as write_wav
|
|
from utils.prompt_making import make_prompt
|
|
|
|
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
|
|
|
|
# Implement the BackendServicer class with the service methods
|
|
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):
|
|
model_name = request.Model
|
|
try:
|
|
print("Preparing models, please wait", file=sys.stderr)
|
|
# download and load all models
|
|
preload_models()
|
|
# Assume directory from request.ModelFile.
|
|
# Only if request.LoraAdapter it's not an absolute path
|
|
if 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
|
|
request.AudioPath = os.path.join(modelFileBase, request.AudioPath)
|
|
if request.AudioPath != "":
|
|
print("Generating model", file=sys.stderr)
|
|
make_prompt(name=model_name, audio_prompt_path=request.AudioPath)
|
|
### Use given transcript
|
|
##make_prompt(name=model_name, audio_prompt_path="paimon_prompt.wav",
|
|
## transcript="Just, what was that? Paimon thought we were gonna get eaten.")
|
|
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):
|
|
model = request.model
|
|
print(request, file=sys.stderr)
|
|
try:
|
|
audio_array = None
|
|
if model != "":
|
|
audio_array = generate_audio(request.text, prompt=model)
|
|
else:
|
|
audio_array = generate_audio(request.text)
|
|
print("saving to", request.dst, file=sys.stderr)
|
|
# save audio to disk
|
|
write_wav(request.dst, SAMPLE_RATE, audio_array)
|
|
print("saved to", request.dst, file=sys.stderr)
|
|
print("tts for", file=sys.stderr)
|
|
print(request, file=sys.stderr)
|
|
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=1))
|
|
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) |