323 lines
12 KiB
Python
323 lines
12 KiB
Python
from typing import Annotated, Union
|
||
from fastapi import FastAPI, Header, status
|
||
from fastapi.responses import Response
|
||
from pydantic import BaseModel
|
||
from enum import Enum
|
||
import os
|
||
import time
|
||
import jwt
|
||
import json
|
||
import base64
|
||
import click
|
||
import ffmpeg
|
||
import shutil
|
||
import logging
|
||
import uvicorn
|
||
import uuid
|
||
from config import settings
|
||
os.environ["TF_CPP_MIN_LOG_LEVEL"] = settings["TF_CPP_MIN_LOG_LEVEL"]
|
||
from spleeter.separator import Separator
|
||
|
||
|
||
class Payload(BaseModel):
|
||
filename: str
|
||
data: str # base64 encoded
|
||
actions: list[str]
|
||
|
||
|
||
class TokenResponse(Enum):
|
||
TOKEN_OK = 0
|
||
TOKEN_SCHEME_ERROR = 1
|
||
TOKEN_EXPIRED = 2
|
||
TOKEN_INVALID = 3
|
||
TOKEN_DECODE_ERROR = 4
|
||
TOKEN_USED = 5
|
||
|
||
class ProcessStatus():
|
||
uuid = ""
|
||
name = ""
|
||
status = 0
|
||
total = 0
|
||
tempfiles = []
|
||
|
||
def __init__(self, uuid, name, total):
|
||
self.uuid = uuid
|
||
self.name = name
|
||
self.status = 0
|
||
self.total = total
|
||
|
||
def update(self):
|
||
self.status += 1
|
||
|
||
def update_temp_files(self, tempfile):
|
||
self.tempfiles.append(tempfile)
|
||
|
||
def is_done(self):
|
||
return self.status == self.total
|
||
|
||
def get_status(self):
|
||
return self.status, self.total
|
||
|
||
def __eq__(self, other):
|
||
return self.uuid == other.uuid or self.name == other.name
|
||
|
||
def __dict__(self):
|
||
return {"uuid": self.uuid, "name": self.name, "status": self.status, "total": self.total, "tempfiles": self.tempfiles}
|
||
|
||
|
||
queue_status = []
|
||
|
||
tempfiles = []
|
||
|
||
app = FastAPI(docs_url=None, redoc_url=None)
|
||
|
||
headers = {
|
||
"typ": "JWT",
|
||
"alg": "HS256"
|
||
}
|
||
|
||
usedJTI = set() # Avoid replay attacks
|
||
separator = Separator("spleeter:2stems")
|
||
separator_lock = False
|
||
LOGFORMAT = "[%(asctime)s] [%(levelname)s] <%(threadName)s> %(message)s"
|
||
logging.basicConfig(format=LOGFORMAT, level=settings.LOG_LEVEL)
|
||
|
||
|
||
def verify_token(auth):
|
||
scheme = auth.split()[0]
|
||
if scheme != "Bearer":
|
||
return TokenResponse.TOKEN_SCHEME_ERROR
|
||
token = auth.split()[1]
|
||
try:
|
||
result = jwt.decode(token,
|
||
settings["JWT_SECRET"],
|
||
verify=True,
|
||
algorithms=["HS256"],
|
||
audience="galaxymusichall.convert")
|
||
if result["jti"] in usedJTI:
|
||
return TokenResponse.TOKEN_USED
|
||
usedJTI.add(result["jti"])
|
||
return TokenResponse.TOKEN_OK
|
||
except jwt.ExpiredSignatureError:
|
||
return TokenResponse.TOKEN_EXPIRED
|
||
except jwt.InvalidTokenError:
|
||
return TokenResponse.TOKEN_INVALID
|
||
except jwt.DecodeError:
|
||
return TokenResponse.TOKEN_DECODE_ERROR
|
||
|
||
|
||
@app.post("/clear")
|
||
def clear(authorization: Annotated[Union[str, None], Header()],
|
||
response: Response):
|
||
if verify := verify_token(authorization):
|
||
if verify == TokenResponse.TOKEN_OK:
|
||
pass
|
||
elif verify == TokenResponse.TOKEN_SCHEME_ERROR:
|
||
response.status_code = status.HTTP_400_BAD_REQUEST
|
||
return {"error": "Invalid authorization header, 'bearer' required"}
|
||
elif verify == TokenResponse.TOKEN_EXPIRED:
|
||
response.status_code = status.HTTP_403_FORBIDDEN
|
||
return {"error": "Token has expired"}
|
||
elif verify == TokenResponse.TOKEN_INVALID:
|
||
response.status_code = status.HTTP_403_FORBIDDEN
|
||
return {"error": "Invalid token"}
|
||
elif verify == TokenResponse.TOKEN_DECODE_ERROR:
|
||
response.status_code = status.HTTP_400_BAD_REQUEST
|
||
return {"error": "Token decoding failed"}
|
||
elif verify == TokenResponse.TOKEN_USED:
|
||
response.status_code = status.HTTP_403_FORBIDDEN
|
||
return {"error": "Token has been used"}
|
||
else:
|
||
response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
|
||
return {"error": "Uncaught error"}
|
||
for i in tempfiles:
|
||
try:
|
||
os.remove(i)
|
||
logging.info("删除文件: "+i)
|
||
except FileNotFoundError:
|
||
logging.warning("文件已被删除: "+i)
|
||
pass
|
||
except Exception as e:
|
||
logging.error("删除文件失败: "+str(e))
|
||
response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
|
||
return {"error": "Failed to clear temporary files"}
|
||
tempfiles.clear()
|
||
return {"message": "Temporary files cleared"}
|
||
|
||
|
||
@app.get("/status")
|
||
def status(authorization: Annotated[Union[str, None], Header()]):
|
||
if verify := verify_token(authorization):
|
||
if verify == TokenResponse.TOKEN_OK:
|
||
pass
|
||
elif verify == TokenResponse.TOKEN_SCHEME_ERROR:
|
||
response.status_code = status.HTTP_400_BAD_REQUEST
|
||
return {"error": "Invalid authorization header, 'bearer' required"}
|
||
elif verify == TokenResponse.TOKEN_EXPIRED:
|
||
response.status_code = status.HTTP_403_FORBIDDEN
|
||
return {"error": "Token has expired"}
|
||
elif verify == TokenResponse.TOKEN_INVALID:
|
||
response.status_code = status.HTTP_403_FORBIDDEN
|
||
return {"error": "Invalid token"}
|
||
elif verify == TokenResponse.TOKEN_DECODE_ERROR:
|
||
response.status_code = status.HTTP_400_BAD_REQUEST
|
||
return {"error": "Token decoding failed"}
|
||
elif verify == TokenResponse.TOKEN_USED:
|
||
response.status_code = status.HTTP_403_FORBIDDEN
|
||
return {"error": "Token has been used"}
|
||
else:
|
||
response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
|
||
return {"error": "Uncaught error"}
|
||
return json.dumps([i.__dict__() for i in queue_status])
|
||
|
||
@app.post("/process")
|
||
def process(data: Payload,
|
||
authorization: Annotated[Union[str, None], Header()],
|
||
response: Response):
|
||
global separator_lock
|
||
#print(data.filename, data.actions, authorization)
|
||
logging.info(f"处理文件{data.filename}，操作{data.actions}")
|
||
processUUID = str(uuid.uuid4())
|
||
process_status = ProcessStatus(uuid=processUUID, name=data.filename, total=len(data.actions))
|
||
queue_status.append(process_status)
|
||
if verify := verify_token(authorization):
|
||
if verify == TokenResponse.TOKEN_OK:
|
||
pass
|
||
elif verify == TokenResponse.TOKEN_SCHEME_ERROR:
|
||
response.status_code = status.HTTP_400_BAD_REQUEST
|
||
return {"error": "Invalid authorization header, 'bearer' required"}
|
||
elif verify == TokenResponse.TOKEN_EXPIRED:
|
||
response.status_code = status.HTTP_403_FORBIDDEN
|
||
return {"error": "Token has expired"}
|
||
elif verify == TokenResponse.TOKEN_INVALID:
|
||
response.status_code = status.HTTP_403_FORBIDDEN
|
||
return {"error": "Invalid token"}
|
||
elif verify == TokenResponse.TOKEN_DECODE_ERROR:
|
||
response.status_code = status.HTTP_400_BAD_REQUEST
|
||
return {"error": "Token decoding failed"}
|
||
elif verify == TokenResponse.TOKEN_USED:
|
||
response.status_code = status.HTTP_403_FORBIDDEN
|
||
return {"error": "Token has been used"}
|
||
else:
|
||
response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
|
||
return {"error": "Uncaught error"}
|
||
tempfilename = "/tmp/"+data.filename
|
||
try:
|
||
with open(tempfilename, "wb") as wrfile:
|
||
wrfile.write(base64.b64decode(data.data))
|
||
except FileExistsError:
|
||
response.status_code = status.HTTP_409_CONFLICT
|
||
return {"error": "File already exists"}
|
||
except Exception as e:
|
||
response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
|
||
logging.error("Unexpected exception:",e)
|
||
return {"error": "Uncaught error"}
|
||
current_name = tempfilename
|
||
process_status.update_temp_files(tempfilename)
|
||
basename = tempfilename[:tempfilename.rfind(".")]
|
||
try:
|
||
for i in data.actions:
|
||
logging.info(f"正在处理{process_status.get_status()}/{process_status.total}：{i}")
|
||
if i == "separate":
|
||
# TODO:性能优化，并发
|
||
while separator_lock:
|
||
logging.info("等待上一Spleeter任务完成...")
|
||
time.sleep(2)
|
||
separator_lock = True
|
||
current_name = separate(current_name, basename)
|
||
process_status.update_temp_files(basename+"/vocals.wav")
|
||
process_status.update_temp_files(current_name)
|
||
separator_lock = False
|
||
elif i == "audio":
|
||
current_name = audio(current_name, basename, target="aac")
|
||
process_status.update_temp_files(current_name)
|
||
elif "convert" in i:
|
||
if "->" not in i:
|
||
response.status_code = status.HTTP_400_BAD_REQUEST
|
||
return {"error": "操作格式错误", "hint": "使用'->'指定目标格式"}
|
||
current_name = convert(current_name, basename, target=i.split("->")[1])
|
||
process_status.update_temp_files(current_name)
|
||
else:
|
||
logging.warning("无效操作"+i+"，跳过中...")
|
||
pass
|
||
process_status.update()
|
||
if current_name is None:
|
||
response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
|
||
return {"error": "处理失败", "status": process_status.get_status(), "required": process_status.total}
|
||
data = {"filename": current_name[current_name.rfind('/') + 1:],
|
||
"data": base64.b64encode(open(current_name, "rb").read()).decode()}
|
||
if process_status.is_done():
|
||
tempfiles.extend(process_status.tempfiles)
|
||
queue_status.remove(process_status)
|
||
else:
|
||
response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
|
||
return {"error": "未预期的错误：流水线在完成前就退出", "status": process_status.get_status(), "required": process_status.total}
|
||
return json.dumps(data)
|
||
except Exception as e:
|
||
logging.error("未预期的错误：" + str(e))
|
||
response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
|
||
return {"error": "Processing failed"}
|
||
|
||
|
||
def separate(filename, basename):
|
||
separator.separate_to_file(filename, "/tmp")
|
||
shutil.move(basename + "/accompaniment.wav", basename + ".wav")
|
||
return basename + ".wav"
|
||
|
||
|
||
def audio(filename, basename, target="aac"):
|
||
audio = ffmpeg.input(filename).audio
|
||
out = ffmpeg.output(audio, basename+"."+target)
|
||
out.overwrite_output().run()
|
||
return basename+"."+target
|
||
|
||
|
||
def convert(filename, basename, target="wav"):
|
||
audio = ffmpeg.input(filename)
|
||
out = ffmpeg.output(audio, basename+"."+target)
|
||
out.overwrite_output().run()
|
||
return basename+"."+target
|
||
|
||
|
||
@click.group()
|
||
def cli():
|
||
pass
|
||
|
||
|
||
@cli.command()
|
||
def secret():
|
||
"""显示JWT密钥"""
|
||
click.echo(settings["JWT_SECRET"])
|
||
|
||
|
||
@cli.command()
|
||
@click.argument("token")
|
||
def verify(token):
|
||
"""验证JWT token"""
|
||
try:
|
||
result = jwt.decode(token,
|
||
settings["JWT_SECRET"],
|
||
verify=True,
|
||
algorithms=["HS256"],
|
||
audience="galaxymusichall.convert")
|
||
click.echo(result)
|
||
except jwt.DecodeError as e:
|
||
click.echo("JWT token解码失败"+str(e))
|
||
except jwt.ExpiredSignatureError:
|
||
click.echo("token已过期。")
|
||
except jwt.InvalidTokenError:
|
||
click.echo("token无效。")
|
||
except Exception as e:
|
||
click.echo("未知错误：")
|
||
click.echo(e)
|
||
|
||
|
||
@cli.command()
|
||
def run():
|
||
"""启动转换服务器"""
|
||
uvicorn.run(app, host=settings.HOST, port=settings.PORT)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
cli()
|