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 "")