404 lines
14 KiB
Python
404 lines
14 KiB
Python
import glob
|
||
import os
|
||
import pathlib
|
||
import shutil
|
||
from typing import Union
|
||
|
||
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
|
||
from app.controllers import base
|
||
from app.controllers.manager.base_manager import TaskQueueFullError
|
||
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 (
|
||
AudioRequest,
|
||
BgmRetrieveResponse,
|
||
BgmUploadResponse,
|
||
SubtitleRequest,
|
||
TaskDeletionResponse,
|
||
TaskQueryRequest,
|
||
TaskQueryResponse,
|
||
TaskResponse,
|
||
TaskVideoRequest,
|
||
VideoMaterialUploadResponse,
|
||
VideoMaterialRetrieveResponse
|
||
)
|
||
from app.services import state as sm
|
||
from app.services import task as tm
|
||
from app.utils import file_security, utils
|
||
|
||
# 认证依赖项
|
||
# router = new_router(dependencies=[Depends(base.verify_token)])
|
||
router = new_router()
|
||
|
||
_enable_redis = config.app.get("enable_redis", False)
|
||
_redis_host = config.app.get("redis_host", "localhost")
|
||
_redis_port = config.app.get("redis_port", 6379)
|
||
_redis_db = config.app.get("redis_db", 0)
|
||
_redis_password = config.app.get("redis_password", None)
|
||
_max_concurrent_tasks = config.app.get("max_concurrent_tasks", 5)
|
||
_max_queued_tasks = config.app.get("max_queued_tasks", 100)
|
||
|
||
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,
|
||
max_queued_tasks=_max_queued_tasks,
|
||
)
|
||
else:
|
||
task_manager = InMemoryTaskManager(
|
||
max_concurrent_tasks=_max_concurrent_tasks,
|
||
max_queued_tasks=_max_queued_tasks,
|
||
)
|
||
|
||
|
||
def _sanitize_upload_filename(filename: str, request_id: str) -> str:
|
||
# 浏览器或客户端有时会附带目录信息，甚至可能夹带 ../ 这类穿越片段。
|
||
# 这里只保留纯文件名，避免上传接口把文件写到目标目录之外。
|
||
normalized_name = (filename or "").replace("\\", "/").split("/")[-1].strip()
|
||
if not normalized_name or normalized_name in {".", ".."}:
|
||
raise HttpException(
|
||
task_id=request_id,
|
||
status_code=400,
|
||
message=f"{request_id}: invalid filename",
|
||
)
|
||
return normalized_name
|
||
|
||
|
||
def _resolve_path_within_directory(base_dir: str, unsafe_path: str, request_id: str) -> str:
|
||
try:
|
||
return file_security.resolve_path_within_directory(base_dir, unsafe_path)
|
||
except ValueError as exc:
|
||
logger.warning(
|
||
f"reject unsafe file path, request_id: {request_id}, path: {unsafe_path}, "
|
||
f"error: {str(exc)}"
|
||
)
|
||
raise HttpException(
|
||
task_id=request_id,
|
||
status_code=404 if str(exc) == "file does not exist" else 403,
|
||
message=f"{request_id}: invalid file path",
|
||
)
|
||
|
||
def _task_file_to_uri(file: str, endpoint: str, task_dir: str, request_id: str) -> str:
|
||
if not isinstance(file, str):
|
||
return file
|
||
|
||
if file.startswith(("http://", "https://")):
|
||
return file
|
||
|
||
try:
|
||
resolved_path = file_security.resolve_path_within_directory(task_dir, file)
|
||
except ValueError as exc:
|
||
# 任务状态理论上只应保存任务目录内的产物路径。这里不再继续拼接 URL，
|
||
# 避免把异常路径包装成可访问链接；同时保留原值，便于排查历史脏数据。
|
||
logger.warning(
|
||
f"skip unsafe task output path, request_id: {request_id}, path: {file}, "
|
||
f"error: {str(exc)}"
|
||
)
|
||
return file
|
||
|
||
relative_path = os.path.relpath(resolved_path, task_dir).replace("\\", "/")
|
||
uri_path = f"tasks/{relative_path}"
|
||
if endpoint:
|
||
return f"{endpoint.rstrip('/')}/{uri_path}"
|
||
return f"/{uri_path}"
|
||
|
||
|
||
@router.post("/videos", response_model=TaskResponse, summary="Generate a short video")
|
||
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.model_dump(),
|
||
}
|
||
sm.state.update_task(task_id)
|
||
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 TaskQueueFullError as e:
|
||
sm.state.delete_task(task_id)
|
||
logger.warning(
|
||
f"reject task because queue is full, request_id: {request_id}, task_id: {task_id}"
|
||
)
|
||
raise HttpException(
|
||
task_id=task_id, status_code=429, message=f"{request_id}: {str(e)}"
|
||
)
|
||
except ValueError as e:
|
||
raise HttpException(
|
||
task_id=task_id, status_code=400, message=f"{request_id}: {str(e)}"
|
||
)
|
||
|
||
from fastapi import Query
|
||
|
||
@router.get("/tasks", response_model=TaskQueryResponse, summary="Get all tasks")
|
||
def get_all_tasks(request: Request, page: int = Query(1, ge=1), page_size: int = Query(10, ge=1)):
|
||
request_id = base.get_task_id(request)
|
||
tasks, total = sm.state.get_all_tasks(page, page_size)
|
||
|
||
response = {
|
||
"tasks": tasks,
|
||
"total": total,
|
||
"page": page,
|
||
"page_size": page_size,
|
||
}
|
||
return utils.get_response(200, response)
|
||
|
||
|
||
|
||
@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(),
|
||
):
|
||
request_id = base.get_task_id(request)
|
||
endpoint = config.app.get("endpoint", "").rstrip("/")
|
||
task = sm.state.get_task(task_id)
|
||
if task:
|
||
task_dir = utils.task_dir()
|
||
response_task = dict(task)
|
||
|
||
if "videos" in task:
|
||
response_task["videos"] = [
|
||
_task_file_to_uri(v, endpoint, task_dir, request_id)
|
||
for v in task["videos"]
|
||
]
|
||
if "combined_videos" in task:
|
||
response_task["combined_videos"] = [
|
||
_task_file_to_uri(v, endpoint, task_dir, request_id)
|
||
for v in task["combined_videos"]
|
||
]
|
||
return utils.get_response(200, response_task)
|
||
|
||
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",
|
||
)
|
||
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)
|
||
if task:
|
||
tasks_dir = utils.task_dir()
|
||
current_task_dir = os.path.join(tasks_dir, task_id)
|
||
if os.path.exists(current_task_dir):
|
||
shutil.rmtree(current_task_dir)
|
||
|
||
sm.state.delete_task(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"
|
||
)
|
||
|
||
|
||
@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:
|
||
filename = os.path.basename(file)
|
||
bgm_list.append(
|
||
{
|
||
"name": filename,
|
||
"size": os.path.getsize(file),
|
||
# 只返回文件名，避免把服务器绝对路径暴露给调用方。
|
||
# 服务端后续会把该文件名解析回 songs 白名单目录。
|
||
"file": filename,
|
||
}
|
||
)
|
||
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",
|
||
)
|
||
def upload_bgm_file(request: Request, file: UploadFile = File(...)):
|
||
request_id = base.get_task_id(request)
|
||
safe_filename = _sanitize_upload_filename(file.filename, request_id)
|
||
# check file ext
|
||
if safe_filename.lower().endswith("mp3"):
|
||
song_dir = utils.song_dir()
|
||
save_path = os.path.join(song_dir, safe_filename)
|
||
# save file
|
||
with open(save_path, "wb+") as buffer:
|
||
# If the file already exists, it will be overwritten
|
||
file.file.seek(0)
|
||
buffer.write(file.file.read())
|
||
response = {"file": safe_filename}
|
||
return utils.get_response(200, response)
|
||
|
||
raise HttpException(
|
||
"", status_code=400, message=f"{request_id}: Only *.mp3 files can be uploaded"
|
||
)
|
||
|
||
@router.get(
|
||
"/video_materials", response_model=VideoMaterialRetrieveResponse, summary="Retrieve local video materials"
|
||
)
|
||
def get_video_materials_list(request: Request):
|
||
allowed_suffixes = ("mp4", "mov", "avi", "flv", "mkv", "jpg", "jpeg", "png")
|
||
local_videos_dir = utils.storage_dir("local_videos", create=True)
|
||
files = []
|
||
for suffix in allowed_suffixes:
|
||
files.extend(glob.glob(os.path.join(local_videos_dir, f"*.{suffix}")))
|
||
# 文件系统枚举顺序不稳定，直接返回会导致“顺序拼接”在不同机器或不同
|
||
# 时刻表现不一致。这里统一按文件名排序，至少保证服务端返回顺序可预测。
|
||
files.sort(key=lambda file_path: os.path.basename(file_path).lower())
|
||
video_materials_list = []
|
||
for file in files:
|
||
filename = os.path.basename(file)
|
||
video_materials_list.append(
|
||
{
|
||
"name": filename,
|
||
"size": os.path.getsize(file),
|
||
# 与 BGM 一样，只返回文件名；创建任务时再在 local_videos
|
||
# 白名单目录内解析，避免 API 泄露宿主机绝对路径。
|
||
"file": filename,
|
||
}
|
||
)
|
||
response = {"files": video_materials_list}
|
||
return utils.get_response(200, response)
|
||
|
||
|
||
@router.post(
|
||
"/video_materials",
|
||
response_model=VideoMaterialUploadResponse,
|
||
summary="Upload the video material file to the local videos directory",
|
||
)
|
||
def upload_video_material_file(request: Request, file: UploadFile = File(...)):
|
||
request_id = base.get_task_id(request)
|
||
safe_filename = _sanitize_upload_filename(file.filename, request_id)
|
||
# check file ext
|
||
allowed_suffixes = ("mp4", "mov", "avi", "flv", "mkv", "jpg", "jpeg", "png")
|
||
normalized_filename = safe_filename.lower()
|
||
# 统一按小写扩展名校验，兼容 .MOV 这类大写后缀文件。
|
||
if normalized_filename.endswith(allowed_suffixes):
|
||
local_videos_dir = utils.storage_dir("local_videos", create=True)
|
||
save_path = os.path.join(local_videos_dir, safe_filename)
|
||
# save file
|
||
with open(save_path, "wb+") as buffer:
|
||
# If the file already exists, it will be overwritten
|
||
file.file.seek(0)
|
||
buffer.write(file.file.read())
|
||
response = {"file": safe_filename}
|
||
return utils.get_response(200, response)
|
||
|
||
raise HttpException(
|
||
"", status_code=400, message=f"{request_id}: Only files with extensions {', '.join(allowed_suffixes)} can be uploaded"
|
||
)
|
||
|
||
@router.get("/stream/{file_path:path}")
|
||
async def stream_video(request: Request, file_path: str):
|
||
request_id = base.get_task_id(request)
|
||
tasks_dir = utils.task_dir()
|
||
video_path = _resolve_path_within_directory(tasks_dir, file_path, request_id)
|
||
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("-")]
|
||
if start is None:
|
||
start = video_size - end
|
||
end = video_size - 1
|
||
if end is None:
|
||
end = video_size - 1
|
||
length = end - start + 1
|
||
|
||
def file_iterator(file_path, offset=0, bytes_to_read=None):
|
||
with open(file_path, "rb") as f:
|
||
f.seek(offset, os.SEEK_SET)
|
||
remaining = bytes_to_read or video_size
|
||
while remaining > 0:
|
||
bytes_to_read = min(4096, remaining)
|
||
data = f.read(bytes_to_read)
|
||
if not data:
|
||
break
|
||
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.status_code = 206 # Partial Content
|
||
|
||
return response
|
||
|
||
|
||
@router.get("/download/{file_path:path}")
|
||
async def download_video(request: Request, file_path: str):
|
||
"""
|
||
download video
|
||
:param request: Request request
|
||
:param file_path: video file path, eg: /cd1727ed-3473-42a2-a7da-4faafafec72b/final-1.mp4
|
||
:return: video file
|
||
"""
|
||
request_id = base.get_task_id(request)
|
||
tasks_dir = utils.task_dir()
|
||
video_path = _resolve_path_within_directory(tasks_dir, file_path, request_id)
|
||
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:]}",
|
||
)
|