Skip to content

Commit

Permalink
feat: add websocket endpoint for fal
Browse files Browse the repository at this point in the history
  • Loading branch information
kennethnym committed Jul 25, 2024
1 parent 0e744af commit 37cf800
Showing 1 changed file with 14 additions and 0 deletions.
14 changes: 14 additions & 0 deletions fal_app.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import io
import fal
from fastapi import WebSocket
import torch
from fal.toolkit import File

Expand Down Expand Up @@ -33,3 +34,16 @@ def run(self):
serialized.append(buf.getvalue())

return serialized

@fal.endpoint("/ws")
async def run_ws(self, ws: WebSocket):
await ws.accept()

wav = self.model.generate(PROMPTS)

for one_wav in enumerate(wav):
buf = io.BytesIO()
torch.save(one_wav, buf)
await ws.send_bytes(buf.getvalue())

await ws.close()

0 comments on commit 37cf800

Please sign in to comment.