hzhang/MoneyPrinterTurbo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

1.Refactor task.py to encapsulate separable functions.

Browse Source 
2.Add a new subtitle API.
This commit is contained in:
yyhhyyyyyy
2024-07-23 17:00:23 +08:00
parent 6d520a4266
commit 17b4a61e64
3 changed files with 383 additions and 194 deletions
Download Patch File Download Diff File Expand all files Collapse all files

169
app/controllers/v1/video.py
View File

@@ -1,11 +1,12 @@
import os
import glob
import os
import pathlib
import shutil
from typing import Union
from fastapi import Request, Depends, Path, BackgroundTasks, UploadFile
from fastapi.responses import FileResponse, StreamingResponse
from fastapi import BackgroundTasks, Depends, Path, Request, UploadFile
from fastapi.params import File
from fastapi.responses import FileResponse, StreamingResponse
from loguru import logger
from app.config import config
@@ -14,10 +15,19 @@ from app.controllers.manager.memory_manager import InMemoryTaskManager
from app.controllers.manager.redis_manager import RedisTaskManager
from app.controllers.v1.base import new_router
from app.models.exception import HttpException
from app.models.schema import TaskVideoRequest, TaskQueryResponse, TaskResponse, TaskQueryRequest, \
BgmUploadResponse, BgmRetrieveResponse, TaskDeletionResponse
from app.services import task as tm
from app.models.schema import (
AudioRequest,
BgmRetrieveResponse,
BgmUploadResponse,
SubtitleRequest,
TaskDeletionResponse,
TaskQueryRequest,
TaskQueryResponse,
TaskResponse,
TaskVideoRequest,
)
from app.services import state as sm
from app.services import task as tm
from app.utils import utils
# 认证依赖项
@@ -34,48 +44,65 @@ _max_concurrent_tasks = config.app.get("max_concurrent_tasks", 5)
redis_url = f"redis://:{_redis_password}@{_redis_host}:{_redis_port}/{_redis_db}"
# 根据配置选择合适的任务管理器
if _enable_redis:
task_manager = RedisTaskManager(max_concurrent_tasks=_max_concurrent_tasks, redis_url=redis_url)
task_manager = RedisTaskManager(
max_concurrent_tasks=_max_concurrent_tasks, redis_url=redis_url
)
else:
task_manager = InMemoryTaskManager(max_concurrent_tasks=_max_concurrent_tasks)
# @router.post("/videos-test", response_model=TaskResponse, summary="Generate a short video")
# async def create_video_test(request: Request, body: TaskVideoRequest):
# task_id = utils.get_uuid()
# request_id = base.get_task_id(request)
# try:
# task = {
# "task_id": task_id,
# "request_id": request_id,
# "params": body.dict(),
# }
# task_manager.add_task(tm.start_test, task_id=task_id, params=body)
# return utils.get_response(200, task)
# except ValueError as e:
# raise HttpException(task_id=task_id, status_code=400, message=f"{request_id}: {str(e)}")
@router.post("/videos", response_model=TaskResponse, summary="Generate a short video")
def create_video(background_tasks: BackgroundTasks, request: Request, body: TaskVideoRequest):
def create_video(
background_tasks: BackgroundTasks, request: Request, body: TaskVideoRequest
):
return create_task(request, body, stop_at="video")
@router.post("/subtitle", response_model=TaskResponse, summary="Generate subtitle only")
def create_subtitle(
background_tasks: BackgroundTasks, request: Request, body: SubtitleRequest
):
return create_task(request, body, stop_at="subtitle")
@router.post("/audio", response_model=TaskResponse, summary="Generate audio only")
def create_audio(
background_tasks: BackgroundTasks, request: Request, body: AudioRequest
):
return create_task(request, body, stop_at="audio")
def create_task(
request: Request,
body: Union[TaskVideoRequest, SubtitleRequest, AudioRequest],
stop_at: str,
):
task_id = utils.get_uuid()
request_id = base.get_task_id(request)
try:
task = {
"task_id": task_id,
"request_id": request_id,
"params": body.dict(),
"params": body.model_dump(),
}
sm.state.update_task(task_id)
# background_tasks.add_task(tm.start, task_id=task_id, params=body)
task_manager.add_task(tm.start, task_id=task_id, params=body)
logger.success(f"video created: {utils.to_json(task)}")
task_manager.add_task(tm.start, task_id=task_id, params=body, stop_at=stop_at)
logger.success(f"Task created: {utils.to_json(task)}")
return utils.get_response(200, task)
except ValueError as e:
raise HttpException(task_id=task_id, status_code=400, message=f"{request_id}: {str(e)}")
raise HttpException(
task_id=task_id, status_code=400, message=f"{request_id}: {str(e)}"
)
@router.get("/tasks/{task_id}", response_model=TaskQueryResponse, summary="Query task status")
def get_task(request: Request, task_id: str = Path(..., description="Task ID"),
query: TaskQueryRequest = Depends()):
@router.get(
"/tasks/{task_id}", response_model=TaskQueryResponse, summary="Query task status"
)
def get_task(
request: Request,
task_id: str = Path(..., description="Task ID"),
query: TaskQueryRequest = Depends(),
):
endpoint = config.app.get("endpoint", "")
if not endpoint:
endpoint = str(request.base_url)
@@ -108,10 +135,16 @@ def get_task(request: Request, task_id: str = Path(..., description="Task ID"),
task["combined_videos"] = urls
return utils.get_response(200, task)
raise HttpException(task_id=task_id, status_code=404, message=f"{request_id}: task not found")
raise HttpException(
task_id=task_id, status_code=404, message=f"{request_id}: task not found"
)
@router.delete("/tasks/{task_id}", response_model=TaskDeletionResponse, summary="Delete a generated short video task")
@router.delete(
"/tasks/{task_id}",
response_model=TaskDeletionResponse,
summary="Delete a generated short video task",
)
def delete_video(request: Request, task_id: str = Path(..., description="Task ID")):
request_id = base.get_task_id(request)
task = sm.state.get_task(task_id)
@@ -125,32 +158,40 @@ def delete_video(request: Request, task_id: str = Path(..., description="Task ID
logger.success(f"video deleted: {utils.to_json(task)}")
return utils.get_response(200)
raise HttpException(task_id=task_id, status_code=404, message=f"{request_id}: task not found")
raise HttpException(
task_id=task_id, status_code=404, message=f"{request_id}: task not found"
)
@router.get("/musics", response_model=BgmRetrieveResponse, summary="Retrieve local BGM files")
@router.get(
"/musics", response_model=BgmRetrieveResponse, summary="Retrieve local BGM files"
)
def get_bgm_list(request: Request):
suffix = "*.mp3"
song_dir = utils.song_dir()
files = glob.glob(os.path.join(song_dir, suffix))
bgm_list = []
for file in files:
bgm_list.append({
"name": os.path.basename(file),
"size": os.path.getsize(file),
"file": file,
})
response = {
"files": bgm_list
}
bgm_list.append(
{
"name": os.path.basename(file),
"size": os.path.getsize(file),
"file": file,
}
)
response = {"files": bgm_list}
return utils.get_response(200, response)
@router.post("/musics", response_model=BgmUploadResponse, summary="Upload the BGM file to the songs directory")
@router.post(
"/musics",
response_model=BgmUploadResponse,
summary="Upload the BGM file to the songs directory",
)
def upload_bgm_file(request: Request, file: UploadFile = File(...)):
request_id = base.get_task_id(request)
# check file ext
if file.filename.endswith('mp3'):
if file.filename.endswith("mp3"):
song_dir = utils.song_dir()
save_path = os.path.join(song_dir, file.filename)
# save file
@@ -158,26 +199,26 @@ def upload_bgm_file(request: Request, file: UploadFile = File(...)):
# If the file already exists, it will be overwritten
file.file.seek(0)
buffer.write(file.file.read())
response = {
"file": save_path
}
response = {"file": save_path}
return utils.get_response(200, response)
raise HttpException('', status_code=400, message=f"{request_id}: Only *.mp3 files can be uploaded")
raise HttpException(
"", status_code=400, message=f"{request_id}: Only *.mp3 files can be uploaded"
)
@router.get("/stream/{file_path:path}")
async def stream_video(request: Request, file_path: str):
tasks_dir = utils.task_dir()
video_path = os.path.join(tasks_dir, file_path)
range_header = request.headers.get('Range')
range_header = request.headers.get("Range")
video_size = os.path.getsize(video_path)
start, end = 0, video_size - 1
length = video_size
if range_header:
range_ = range_header.split('bytes=')[1]
start, end = [int(part) if part else None for part in range_.split('-')]
range_ = range_header.split("bytes=")[1]
start, end = [int(part) if part else None for part in range_.split("-")]
if start is None:
start = video_size - end
end = video_size - 1
@@ -186,7 +227,7 @@ async def stream_video(request: Request, file_path: str):
length = end - start + 1
def file_iterator(file_path, offset=0, bytes_to_read=None):
with open(file_path, 'rb') as f:
with open(file_path, "rb") as f:
f.seek(offset, os.SEEK_SET)
remaining = bytes_to_read or video_size
while remaining > 0:
@@ -197,10 +238,12 @@ async def stream_video(request: Request, file_path: str):
remaining -= len(data)
yield data
response = StreamingResponse(file_iterator(video_path, start, length), media_type='video/mp4')
response.headers['Content-Range'] = f'bytes {start}-{end}/{video_size}'
response.headers['Accept-Ranges'] = 'bytes'
response.headers['Content-Length'] = str(length)
response = StreamingResponse(
file_iterator(video_path, start, length), media_type="video/mp4"
)
response.headers["Content-Range"] = f"bytes {start}-{end}/{video_size}"
response.headers["Accept-Ranges"] = "bytes"
response.headers["Content-Length"] = str(length)
response.status_code = 206 # Partial Content
return response
@@ -219,8 +262,10 @@ async def download_video(_: Request, file_path: str):
file_path = pathlib.Path(video_path)
filename = file_path.stem
extension = file_path.suffix
headers = {
"Content-Disposition": f"attachment; filename={filename}{extension}"
}
return FileResponse(path=video_path, headers=headers, filename=f"{filename}{extension}",
media_type=f'video/{extension[1:]}')
headers = {"Content-Disposition": f"attachment; filename={filename}{extension}"}
return FileResponse(
path=video_path,
headers=headers,
filename=f"{filename}{extension}",
media_type=f"video/{extension[1:]}",
)