fix: websocket issue

websocket.create_connection is short lived, so whenever the inference
server is needed, it needs to be called again to reconnect to the
server.
This commit is contained in:
2024-07-23 21:43:35 +01:00
parent 63a8056283
commit a6d6487df3
2 changed files with 13 additions and 9 deletions

View File

@@ -22,7 +22,7 @@ async def handler(websocket):
async def main(): async def main():
async with serve(handler, "", 8001): async with serve(handler, "localhost", 8001, ping_interval=5):
await asyncio.Future() await asyncio.Future()

View File

@@ -13,17 +13,16 @@ current_index = -1
t = None t = None
# websocket connection to the inference server # websocket connection to the inference server
ws = None ws = None
ws_url = ""
@asynccontextmanager @asynccontextmanager
async def lifespan(app: FastAPI): async def lifespan(app: FastAPI):
global ws global ws, ws_url
url = os.environ.get("INFERENCE_SERVER_WS_URL") ws_url = os.environ.get("INFERENCE_SERVER_WS_URL")
if not url: if not ws_url:
url = "ws://localhost:8001" ws_url = "ws://localhost:8001"
ws = websocket.create_connection(url)
print(f"websocket connected to {url}")
advance() advance()
@@ -36,7 +35,7 @@ async def lifespan(app: FastAPI):
def generate_new_audio(): def generate_new_audio():
if not ws: if not ws_url:
return return
global current_index global current_index
@@ -51,6 +50,9 @@ def generate_new_audio():
print("generating new audio...") print("generating new audio...")
ws = websocket.create_connection(ws_url)
print(f"websocket connected to {ws_url}")
ws.send("generate") ws.send("generate")
wavs = [] wavs = []
@@ -66,6 +68,8 @@ def generate_new_audio():
print("audio generated.") print("audio generated.")
ws.close()
def advance(): def advance():
global current_index, t global current_index, t