From c43df2ac723c5b573223e0420785af974f4c1151 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Sat, 2 Sep 2023 23:44:10 +0800
Subject: [PATCH] feat: use av instead of ffmpeg

---
 .gitignore       |  2 ++
 requirements.txt |  2 +-
 server.py        | 44 +++++++++++++++-----------------------------
 utils.py         | 16 ++++++++++++++++
 4 files changed, 34 insertions(+), 30 deletions(-)

diff --git a/.gitignore b/.gitignore
index dc12b7f..a7ecc18 100644
--- a/.gitignore
+++ b/.gitignore
@@ -159,4 +159,6 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 
+.DS_Store
 /models
+/logs
diff --git a/requirements.txt b/requirements.txt
index bfffc61..c437b9a 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,3 @@
-Cython
 librosa==0.9.1
 matplotlib
 numpy
@@ -15,3 +14,4 @@ transformers
 pypinyin
 cn2an
 gradio
+av
diff --git a/server.py b/server.py
index 025731b..552f491 100644
--- a/server.py
+++ b/server.py
@@ -1,28 +1,15 @@
 from flask import Flask, request, Response
 from io import BytesIO
-import ffmpeg
-import base64
-import os
-import sys
-import json
-import math
 import torch
-from torch import nn
-from torch.nn import functional as F
-from torch.utils.data import DataLoader
 
-import time
 import commons
 import utils
 from models import SynthesizerTrn
 from text.symbols import symbols
-from text import cleaned_text_to_sequence,_symbol_to_id, get_bert
+from text import cleaned_text_to_sequence, get_bert
 from text.cleaner import clean_text
 from scipy.io import wavfile
 
-# Get ffmpeg path
-ffmpeg_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ffmpeg")
-
 # Flask Init
 app = Flask(__name__)
 app.config['JSON_AS_ASCII'] = False
@@ -98,24 +85,23 @@ def main():
             fmt = request.args.get("format", "wav")
             if None in (speaker, text):
                 return "Missing Parameter"
-            if fmt not in ("mp3", "wav"):
+            if fmt not in ("mp3", "wav", "ogg"):
                 return "Invalid Format"
         except:
             return "Invalid Parameter"
 
-
         with torch.no_grad():
             audio = infer(text, sdp_ratio=sdp_ratio, noise_scale=noise, noise_scale_w=noisew, length_scale=length, sid=speaker)
-        wav = BytesIO()
-        wavfile.write(wav, hps.data.sampling_rate, audio)
-        torch.cuda.empty_cache()
-        if fmt == "mp3":
-            process = (
-	    ffmpeg
-            .input("pipe:", format='wav', channel_layout="mono")
-            .output("pipe:", format='mp3', audio_bitrate="320k")
-            .run_async(pipe_stdin=True, pipe_stdout=True, pipe_stderr=True, cmd=ffmpeg_path)
-        )
-            out, _ = process.communicate(input=wav.read())
-            return Response(out, mimetype="audio/mpeg")
-        return Response(wav.read(), mimetype="audio/wav")
+
+        with BytesIO() as wav:
+            wavfile.write(wav, hps.data.sampling_rate, audio)
+            torch.cuda.empty_cache()
+            if fmt == "wav":
+                return Response(wav.getvalue(), mimetype="audio/wav")
+            wav.seek(0, 0)
+            with BytesIO() as ofp:
+                utils.wav2(wav, ofp, fmt)
+                return Response(
+                    ofp.getvalue(),
+                    mimetype="audio/mpeg" if fmt == "mp3" else "audio/ogg"
+                )
diff --git a/utils.py b/utils.py
index fbd7149..e917a78 100644
--- a/utils.py
+++ b/utils.py
@@ -8,6 +8,7 @@ import subprocess
 import numpy as np
 from scipy.io.wavfile import read
 import torch
+from av import open as avopen
 
 MATPLOTLIB_FLAG = False
 
@@ -291,3 +292,18 @@ class HParams():
 
     def __repr__(self):
         return self.__dict__.__repr__()
+
+def wav2(i, o, format):
+  inp = avopen(i, 'rb')
+  out = avopen(o, 'wb', format=format)
+  if format == "ogg": format = "libvorbis"
+
+  ostream = out.add_stream(format)
+
+  for frame in inp.decode(audio=0):
+      for p in ostream.encode(frame): out.mux(p)
+
+  for p in ostream.encode(None): out.mux(p)
+
+  out.close()
+  inp.close()