193 lines
8.1 KiB
Python
193 lines
8.1 KiB
Python
import modal
|
|
from typing import Annotated
|
|
from fastapi import FastAPI, Response, Depends, Header, Request
|
|
from starlette import status
|
|
from scalar_fastapi import get_scalar_api_reference
|
|
import sentry_sdk
|
|
from sentry_sdk.integrations.loguru import LoguruIntegration, LoggingLevels
|
|
from sentry_sdk.integrations.fastapi import FastApiIntegration
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from .middleware.authorization import verify_token
|
|
from .utils.KVCache import KVCache
|
|
from .utils.ModalUtils import ModalUtils
|
|
from .models.media_model import MediaSource
|
|
from .models.web_model import (SentryTransactionInfo,
|
|
SentryTransactionHeader,
|
|
ModalTaskResponse,
|
|
ComfyTaskRequest,
|
|
ComfyTaskStatusResponse
|
|
)
|
|
from .router import ffmpeg, cache
|
|
from .config import WorkerConfig
|
|
|
|
config = WorkerConfig()
|
|
|
|
web_app = FastAPI(title="Modal worker API",
|
|
version="0.1.2",
|
|
summary="Modal Worker的API, 包括缓存视频, 发起生产任务等",
|
|
servers=[
|
|
{'url': f'https://bowongai-dev--{config.modal_app_name}-fastapi-webapp.modal.run',
|
|
'description': 'modal 开发环境服务'},
|
|
{'url': 'https://modal-dev.bowong.cc',
|
|
'description': 'modal 开发环境服务独立域名'},
|
|
{'url': f'https://bowongai-test--{config.modal_app_name}-fastapi-webapp.modal.run',
|
|
'description': 'modal 测试环境服务'},
|
|
{'url': f'https://bowongai-main--{config.modal_app_name}-fastapi-webapp.modal.run',
|
|
'description': 'modal 生产环境服务'
|
|
}
|
|
])
|
|
|
|
sentry_sdk.init(dsn="https://dab7b7ae652216282c89f029a76bb10a@sentry.bowongai.com/2",
|
|
send_default_pii=True,
|
|
traces_sample_rate=1.0,
|
|
profiles_sample_rate=1.0,
|
|
add_full_stack=True,
|
|
environment=config.modal_environment,
|
|
integrations=[
|
|
LoguruIntegration(level=LoggingLevels.INFO.value, event_level=LoggingLevels.ERROR.value),
|
|
FastApiIntegration()
|
|
]
|
|
)
|
|
modal_kv_cache = KVCache(kv_name=config.modal_kv_name, environment=config.modal_environment)
|
|
|
|
sentry_header_schema = {
|
|
"x-trace-id": {
|
|
"description": "Sentry Transaction ID",
|
|
"schema": {
|
|
"type": "string",
|
|
}
|
|
},
|
|
"x-baggage": {
|
|
"description": "Sentry Transaction baggage",
|
|
"schema": {
|
|
"type": "string",
|
|
}
|
|
}
|
|
}
|
|
|
|
ALIAS_MAP = {
|
|
"/comfyui": "/comfyui/v2",
|
|
}
|
|
|
|
web_app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
@web_app.middleware("http")
|
|
@web_app.middleware("https")
|
|
async def alias_middleware(request: Request, call_next):
|
|
if request.url.path in ALIAS_MAP.keys():
|
|
request.scope["path"] = ALIAS_MAP[request.url.path]
|
|
return await call_next(request)
|
|
|
|
|
|
@web_app.get("/scalar", include_in_schema=False)
|
|
async def scalar():
|
|
return get_scalar_api_reference(openapi_url=web_app.openapi_schema or '/openapi.json', title="Modal worker web endpoint")
|
|
|
|
|
|
web_app.include_router(ffmpeg.router)
|
|
web_app.include_router(cache.router)
|
|
|
|
|
|
@web_app.post("/comfyui/v1", tags=["ComfyUI"], summary="发起ComfyUIBase任务", description="发起ComfyUIBase任务",
|
|
responses={
|
|
status.HTTP_200_OK: {
|
|
"description": "",
|
|
"headers": sentry_header_schema
|
|
},
|
|
},
|
|
dependencies=[Depends(verify_token)]
|
|
)
|
|
async def comfyui_v1(item: ComfyTaskRequest, headers: Annotated[SentryTransactionHeader, Header()]):
|
|
sentry_trace = None
|
|
if headers.x_trace_id and headers.x_baggage:
|
|
sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
|
|
cls = modal.Cls.from_name(config.modal_app_name, "ComfyUI", environment_name=config.modal_environment)
|
|
fn_call = cls().api.spawn(item.video_path.path, item.start_time, item.filename_prefix, item.tts_text1,
|
|
item.tts_text2, item.tts_text3, item.tts_text4, item.anchor_id, item.speed, sentry_trace,
|
|
item.webhook)
|
|
return ModalTaskResponse(success=True, taskId=fn_call.object_id)
|
|
|
|
|
|
@web_app.get("/comfyui/v1/{task_id}", tags=["ComfyUI"], summary="查询ComfyUIBase任务",
|
|
description="查询ComfyUIBase任务",
|
|
responses={
|
|
status.HTTP_200_OK: {
|
|
"description": "",
|
|
"headers": sentry_header_schema
|
|
},
|
|
},
|
|
dependencies=[Depends(verify_token)]
|
|
)
|
|
async def comfyui_v1_status(task_id: str, response: Response):
|
|
task_status, code, reason, result, transaction = await ModalUtils.get_modal_task_status(task_id)
|
|
media = None
|
|
if transaction:
|
|
response.headers["x-trace-id"] = transaction.x_trace_id
|
|
response.headers["x-baggage"] = transaction.x_baggage
|
|
if task_status == "success":
|
|
task_status = result["status"]
|
|
if task_status != "success":
|
|
reason = result["msg"]
|
|
else:
|
|
media = MediaSource.from_str("s3://" + "/".join(
|
|
[config.S3_region, config.S3_bucket_name, config.comfyui_s3_output, result["file_name"]]))
|
|
return ComfyTaskStatusResponse(taskId=task_id, status=task_status, code=code,
|
|
error=reason, result=media.urn if media else None)
|
|
|
|
|
|
@web_app.post("/comfyui/v2", tags=["ComfyUI"], summary="发起ComfyUILatentSync1.5任务",
|
|
description="发起ComfyUILatentSync1.5任务",
|
|
responses={
|
|
status.HTTP_200_OK: {
|
|
"description": "",
|
|
"headers": sentry_header_schema
|
|
},
|
|
},
|
|
dependencies=[Depends(verify_token)]
|
|
)
|
|
async def comfyui_v2(item: ComfyTaskRequest, headers: Annotated[SentryTransactionHeader, Header()]):
|
|
sentry_trace = None
|
|
if headers.x_trace_id and headers.x_baggage:
|
|
sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
|
|
cls = modal.Cls.from_name(config.modal_app_name, "ComfyUILatentSync15", environment_name=config.modal_environment)
|
|
fn_call = cls().api.spawn(item.video_path.path, item.start_time, item.filename_prefix, item.tts_text1,
|
|
item.tts_text2, item.tts_text3, item.tts_text4, item.anchor_id, item.speed, sentry_trace,
|
|
item.webhook)
|
|
return ModalTaskResponse(success=True, taskId=fn_call.object_id)
|
|
|
|
|
|
@web_app.get("/comfyui/v2/{task_id}", tags=["ComfyUI"], summary="查询ComfyUILatentSync1.5任务",
|
|
description="查询ComfyUILatentSync1.5任务",
|
|
responses={
|
|
status.HTTP_200_OK: {
|
|
"description": "",
|
|
"headers": sentry_header_schema
|
|
},
|
|
},
|
|
dependencies=[Depends(verify_token)]
|
|
)
|
|
async def comfyui_v2_status(task_id: str, response: Response):
|
|
task_status, code, reason, result, transaction = await ModalUtils.get_modal_task_status(task_id)
|
|
media = None
|
|
|
|
if transaction:
|
|
response.headers["x-trace-id"] = transaction.x_trace_id
|
|
response.headers["x-baggage"] = transaction.x_baggage
|
|
if task_status == "success":
|
|
task_status = result["status"]
|
|
if task_status != "success":
|
|
reason = result["msg"]
|
|
else:
|
|
media = MediaSource.from_str("s3://" + "/".join(
|
|
[config.S3_region, config.S3_bucket_name, config.comfyui_s3_output, result["file_name"]]))
|
|
return ComfyTaskStatusResponse(taskId=task_id, status=task_status, code=code,
|
|
error=reason,
|
|
result=media.urn if media else "")
|