Skip to content

Commit

Permalink
fix: write generated audio to fal storage
Browse files Browse the repository at this point in the history
  • Loading branch information
kennethnym committed Aug 20, 2024
1 parent 8e76c91 commit 13ae315
Showing 1 changed file with 17 additions and 5 deletions.
22 changes: 17 additions & 5 deletions fal_app.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,14 @@
import io
from pathlib import Path
from audiocraft.data.audio import audio_write
import fal
from fastapi import WebSocket
import torch
from fal.toolkit import File

from prompts import PROMPTS

DATA_DIR = Path("/data/audio")


class InfinifiFalApp(fal.App, keep_alive=300):
machine_type = "GPU-A6000"
Expand Down Expand Up @@ -41,9 +44,18 @@ async def run_ws(self, ws: WebSocket):

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())
for i, one_wav in enumerate(wav):
path = DATA_DIR.joinpath(f"{i}")
audio_write(
path,
one_wav.cpu(),
self.model.sample_rate,
format="mp3",
strategy="loudness",
loudness_compressor=True,
)
with open(path, "rb") as f:
data = f.read()
await ws.send_bytes(data)

await ws.close()

0 comments on commit 13ae315

Please sign in to comment.