comi-logic/modalDeploy
7
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

ADD AutoDL调度基本功能完成

Browse Source
This commit is contained in:
kyj@bowong.ai
2025-04-16 18:28:15 +08:00
parent 61106608e0
commit 1272a33718
8 changed files with 446 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
AutoDL/AutoDL_pure_heygem.py
View File

@@ -6,12 +6,15 @@ import subprocess
import time
import traceback
import uuid
from concurrent.futures.thread import ThreadPoolExecutor
from typing import Any, Optional, Union
import httpx
import loguru
import requests
import starlette.datastructures
import uvicorn
from fastapi import UploadFile, HTTPException, Depends, FastAPI
from fastapi.routing import APIRoute
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
@@ -47,8 +50,12 @@ class HeyGem:
time.sleep(2)
if timeout == 0:
raise TimeoutError("HeyGem Server Start timed out")
self.result = {}
self.executor = ThreadPoolExecutor(max_workers=1)
self.server = FastAPI()
self.server.routes.append(APIRoute(path='/', endpoint=self.api, methods=['POST']))
self.server.routes.append(APIRoute(path='/{code}', endpoint=self.get_result, methods=['GET']))
def infer(self, code, video_path, audio_path) -> str:
@@ -182,6 +189,18 @@ class HeyGem:
with open(path, 'wb') as f:
shutil.copyfileobj(r.raw, f)
async def get_result(self, code:str, token: HTTPAuthorizationCredentials = Depends(HTTPBearer())):
if token.credentials != os.environ["AUTH_TOKEN"]:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Incorrect bearer token",
headers={"WWW-Authenticate": "Bearer"},
)
if code in self.result:
return self.result[code]
else:
return {"status": "waiting", "msg":""}
async def api(self, video_file: Optional[Union[UploadFile, str]], audio_file: Optional[Union[UploadFile, str]], token: HTTPAuthorizationCredentials = Depends(HTTPBearer())):
if token.credentials != os.environ["AUTH_TOKEN"]:
raise HTTPException(
@@ -190,31 +209,41 @@ class HeyGem:
headers={"WWW-Authenticate": "Bearer"},
)
code = str(uuid.uuid4())
video_path=None
audio_path=None
try:
if isinstance(video_file, UploadFile):
video_path = "/root/%s" % video_file.filename
file_id = str(uuid.uuid4())
if isinstance(video_file, UploadFile) or isinstance(video_file, starlette.datastructures.UploadFile):
video_path = "/root/%s.%s" % (file_id, video_file.filename.split(".")[-1])
with open(video_path, "wb") as f:
data = await video_file.read()
f.write(data)
elif isinstance(video_file, str) and video_file.startswith("http"):
video_path = "/root/%s" % video_file.split("?")[0].split("/")[-1]
video_path = "/root/%s.%s" % (file_id, video_file.split("?")[0].split("/")[-1].split(".")[-1])
self.download_file(video_file, video_path)
else:
return {"status": "fail", "msg": "video file save failed: Not a valid input"}
self.result[code] = {"status": "fail", "msg": "video file save failed: Not a valid input"}
if isinstance(audio_file, UploadFile):
audio_path = "/root/%s" % audio_file.filename
if isinstance(audio_file, UploadFile) or isinstance(audio_file, starlette.datastructures.UploadFile):
audio_path = "/root/%s.%s" % (file_id, audio_file.filename.split(".")[-1])
with open(audio_path, "wb") as f:
data = await audio_file.read()
f.write(data)
elif isinstance(audio_file, str) and audio_file.startswith("http"):
audio_path = "/root/%s" % audio_file.split("?")[0].split("/")[-1]
audio_path = "/root/%s.%s" % (file_id, audio_file.split("?")[0].split("/")[-1].split(".")[-1])
self.download_file(audio_file, audio_path)
else:
return {"status": "fail", "msg": "audio file save failed: Not a valid input"}
self.result[code] = {"status": "fail", "msg": "audio file save failed: Not a valid input"}
except Exception as e:
return {"status": "fail", "msg": f"video or audio file save failed: {str(e)}"}
self.result[code] = {"status": "fail", "msg": f"video or audio file save failed: {str(e)}"}
# Submit
if video_path and audio_path:
self.executor.submit(self._api, video_path=video_path, audio_path=audio_path, code=code)
return {"status": "success", "code": code}
else:
return self.result[code]
def _api(self, video_path: str, audio_path: str, code):
loguru.logger.info(f"Enter Inference Module")
try:
file_path = self.infer(code, video_path, audio_path)
@@ -230,15 +259,11 @@ class HeyGem:
aws_secret_access_key=yaml_config["aws_access_key"])
awss3.meta.client.upload_file(file_path, "bw-heygem-output", file_path.split(os.path.sep)[-1])
except Exception as e:
return {"status": "fail", "msg": "Failed to move file to S3 manually: " + str(e)}
# resp = FileResponse(path=file_path, filename=Path(file_path).name, status_code=200, headers={
# "Content-Type": "application/octet-stream",
# "Content-Disposition": f"attachment; filename={Path(file_path).name}",
# })
return {"status": "success", "msg":f"{file_path.split(os.path.sep)[-1]}"}
self.result[code] = {"status": "fail", "msg": "Failed to move file to S3 manually: " + str(e)}
self.result[code] = {"status": "success", "msg":f"{file_path.split(os.path.sep)[-1]}"}
except Exception as e:
traceback.print_exc()
return {"status": "fail", "msg": "Inference module failed: "+str(e)}
self.result[code] = {"status": "fail", "msg": "Inference module failed: "+str(e)}
if __name__ == "__main__":
heygem = HeyGem()