init: heygen 队列

2025-07-25 15:20:42 +08:00
commit 8301a19003
7 changed files with 1423 additions and 0 deletions
--- a/api.py
+++ b/api.py
@@ -0,0 +1,424 @@
+import asyncio
+import os
+import shutil
+from datetime import datetime
+from typing import Any, Callable, Dict, List, Optional
+
+import aiofiles
+import aiohttp
+import uvicorn
+from fastapi import BackgroundTasks, FastAPI, File, Form, HTTPException, UploadFile
+from fastapi.responses import StreamingResponse
+from loguru import logger
+from models import Task, TaskCreateRequest, TaskStatus
+from processor import ProcessManager
+
+
+# 定义任务管理器
+class TaskManager:
+    def __init__(self):
+        self.tasks: Dict[str, Task] = {}
+        self.process_manager = ProcessManager(config_path="config/machines.json", task_manager=self)
+        self.callbacks: Dict[str, List[Callable[[Task], Any]]] = {}
+        self._task_lock = asyncio.Lock()
+        self._callback_lock = asyncio.Lock()
+        self._started = False
+
+    async def start(self):
+        """启动任务管理器"""
+        if not self._started:
+            await self.process_manager.start()
+            self._started = True
+            logger.info("任务管理器已启动")
+
+    async def stop(self):
+        """停止任务管理器"""
+        if self._started:
+            await self.process_manager.stop()
+            self._started = False
+            logger.info("任务管理器已停止")
+
+    async def add_task(
+        self, audio_file_path: str, video_file_path: str, output_dir: str, description: str, callback_url: Optional[str] = None
+    ) -> Task:
+        """添加新任务"""
+        async with self._task_lock:
+            task_id = f"task_{len(self.tasks) + 1}"
+            output_file_path = os.path.join(output_dir, f"{task_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.mp4")
+            task = Task(
+                task_id=task_id,
+                audio_file_path=audio_file_path,
+                video_file_path=video_file_path,
+                output_dir=output_dir,
+                description=description,
+                status=TaskStatus.PENDING,
+                created_at=datetime.now(),
+                callback_url=callback_url,
+                output_file_path=output_file_path,
+            )
+            self.tasks[task_id] = task
+
+            # 如果有回调URL，添加回调
+            if callback_url:
+                await self.add_callback(task_id, self._create_callback_handler(callback_url))
+
+            return task
+
+    def _create_callback_handler(self, callback_url: str) -> Callable[[Task], Any]:
+        """创建回调处理器"""
+
+        async def callback_handler(task: Task):
+            try:
+                if callback_url.startswith("http"):
+                    async with aiohttp.ClientSession() as session:
+                        await session.post(callback_url, json=task.dict())
+                else:
+                    logger.info(f"回调URL: {callback_url} 不是HTTP URL")
+            except Exception as e:
+                logger.error(f"Error sending callback to {callback_url}: {str(e)}")
+
+        return callback_handler
+
+    async def process_task_sync(self, task_id: str) -> Task:
+        """同步处理任务"""
+        task = await self.get_task(task_id)
+        if not task:
+            raise ValueError(f"Task {task_id} not found")
+
+        try:
+            # 更新任务状态为处理中
+            await self.update_task_status(task_id, TaskStatus.PROCESSING)
+
+            # 生成唯一的文件ID
+            file_id = f"{task_id}_{int(datetime.now().timestamp())}"
+
+            # 处理视频
+            success = await self.process_manager.process_video(
+                task_id=task_id,
+                audio_file_path=task.audio_file_path,
+                video_file_path=task.video_file_path,
+                out_file_path=task.output_file_path,
+            )
+
+            if success:
+                # 更新任务状态为已入队
+                return await self.update_task_status(task_id, TaskStatus.QUEUED)
+            else:
+                return await self.update_task_status(task_id, TaskStatus.FAILED, error_message="Failed to queue video processing task")
+        except Exception as e:
+            return await self.update_task_status(task_id, TaskStatus.FAILED, error_message=str(e))
+
+    async def update_task_status(
+        self, task_id: str, status: TaskStatus, machine_name: Optional[str] = None, error_message: Optional[str] = None
+    ) -> Task:
+        """更新任务状态"""
+        async with self._task_lock:
+            if task_id not in self.tasks:
+                raise ValueError(f"Task {task_id} not found")
+
+            task = self.tasks[task_id]
+            logger.info(f"更新任务状态: {task_id}, 从 {task.status} 到 {status}")
+            task.status = status
+            if machine_name:
+                task.machine_name = machine_name
+            if error_message:
+                task.error_message = error_message
+            if status in [TaskStatus.COMPLETED, TaskStatus.FAILED]:
+                task.completed_at = datetime.now()
+                logger.info(f"任务完成时间已设置: {task_id}, 状态: {status}")
+
+            # 触发回调
+            if task_id in self.callbacks:
+                logger.info(f"触发任务回调: {task_id}, 回调数量: {len(self.callbacks[task_id])}")
+                for callback in self.callbacks[task_id]:
+                    try:
+                        await callback(task)
+                        logger.info(f"任务回调执行成功: {task_id}")
+                    except Exception as e:
+                        logger.error(f"任务回调执行失败: {task_id}, 错误: {str(e)}")
+
+            return task
+
+    async def get_task(self, task_id: str) -> Optional[Task]:
+        """获取任务信息"""
+        return self.tasks.get(task_id)
+
+    async def get_all_tasks(self) -> List[Task]:
+        """获取所有任务"""
+        return list(self.tasks.values())
+
+    async def add_callback(self, task_id: str, callback: Callable[[Task], Any]) -> None:
+        """添加任务回调"""
+        async with self._callback_lock:
+            if task_id not in self.callbacks:
+                self.callbacks[task_id] = []
+            self.callbacks[task_id].append(callback)
+
+    async def remove_callback(self, task_id: str, callback: Callable[[Task], Any]) -> None:
+        """移除任务回调"""
+        async with self._callback_lock:
+            if task_id in self.callbacks:
+                self.callbacks[task_id].remove(callback)
+
+
+# 创建FastAPI应用
+app = FastAPI(title="视频处理API")
+task_manager = TaskManager()
+
+
+@app.on_event("startup")
+async def startup_event():
+    """应用启动时初始化"""
+    await task_manager.start()
+
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    """应用关闭时清理"""
+    await task_manager.stop()
+
+
+# API路由
+@app.post("/tasks/", response_model=Task)
+async def create_task(
+    audio: UploadFile = File(...),
+    video: UploadFile = File(...),
+    description: str = Form(...),
+    callback_url: Optional[str] = Form(None),
+    background_tasks: BackgroundTasks = BackgroundTasks(),
+):
+    """创建新的视频处理任务"""
+    try:
+        # 创建临时目录
+        temp_dir = os.path.join(os.getcwd(), "temp", datetime.now().strftime("%Y%m%d_%H%M%S"))
+        os.makedirs(temp_dir, exist_ok=True)
+        try:
+            # 确保文件名不为None
+            if not audio.filename or not video.filename:
+                raise ValueError("文件名不能为空")
+
+            # 保存上传的文件
+            audio_path = os.path.join(temp_dir, audio.filename)
+            video_path = os.path.join(temp_dir, video.filename)
+
+            # 保存音频文件
+            logger.info(f"开始保存音频文件: {audio_path}")
+            async with aiofiles.open(audio_path, "wb") as f:
+                content = await audio.read()
+                await f.write(content)
+            logger.info(f"音频文件保存完成: {audio_path}")
+
+            # 保存视频文件
+            logger.info(f"开始保存视频文件: {video_path}")
+            async with aiofiles.open(video_path, "wb") as f:
+                content = await video.read()
+                await f.write(content)
+            logger.info(f"视频文件保存完成: {video_path}")
+
+            # 验证文件是否保存成功
+            if not os.path.exists(audio_path) or not os.path.exists(video_path):
+                raise ValueError("文件保存失败")
+
+            # 创建输出目录
+            output_dir = os.path.join("output", datetime.now().strftime("%Y%m%d"))
+            os.makedirs(output_dir, exist_ok=True)
+
+            # 创建任务
+            task = await task_manager.add_task(audio_path, video_path, output_dir, description, callback_url)
+
+            # 异步处理任务
+            background_tasks.add_task(process_video_task, task.task_id, audio_path, video_path, task.output_file_path)
+
+            return task
+
+        except Exception as e:
+            # 如果发生错误，立即清理临时文件
+            # if os.path.exists(temp_dir):
+            #     shutil.rmtree(temp_dir)
+            logger.error(f"创建任务失败: {str(e)}")
+            raise HTTPException(status_code=500, detail=str(e))
+
+    except Exception as e:
+        logger.error(f"创建任务失败: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/tasks/sync")
+async def create_task_sync(
+    audio: UploadFile = File(...),
+    video: UploadFile = File(...),
+    description: str = Form(...),
+    callback_url: Optional[str] = Form(None),
+    background_tasks: BackgroundTasks = BackgroundTasks(),
+):
+    """创建新的视频处理任务并返回处理结果"""
+    try:
+        logger.info("创建新的视频处理任务并返回处理结果")
+        # 创建临时目录
+        temp_dir = os.path.join(os.getcwd(), "temp", datetime.now().strftime("%Y%m%d_%H%M%S"))
+        logger.info(f"临时目录: {temp_dir}")
+        os.makedirs(temp_dir, exist_ok=True)
+        try:
+            # 确保文件名不为None
+            if not audio.filename or not video.filename:
+                raise ValueError("文件名不能为空")
+            logger.info(f"文件名: {audio.filename}, {video.filename}")
+            # 保存上传的文件
+            audio_path = os.path.join(temp_dir, audio.filename)
+            video_path = os.path.join(temp_dir, video.filename)
+
+            # 保存音频文件
+            logger.info(f"开始保存音频文件: {audio_path}")
+            async with aiofiles.open(audio_path, "wb") as f:
+                content = await audio.read()
+                await f.write(content)
+            logger.info(f"音频文件保存完成: {audio_path}")
+
+            # 保存视频文件
+            logger.info(f"开始保存视频文件: {video_path}")
+            async with aiofiles.open(video_path, "wb") as f:
+                content = await video.read()
+                await f.write(content)
+            logger.info(f"视频文件保存完成: {video_path}")
+
+            # 验证文件是否保存成功
+            if not os.path.exists(audio_path) or not os.path.exists(video_path):
+                raise ValueError("文件保存失败")
+
+            # 创建输出目录
+            output_dir = os.path.join("output", datetime.now().strftime("%Y%m%d"))
+            os.makedirs(output_dir, exist_ok=True)
+
+            # 创建任务
+            task = await task_manager.add_task(audio_path, video_path, output_dir, description, callback_url)
+
+            # 创建一个事件来等待任务完成
+            task_completed = asyncio.Event()
+            logger.info(f"创建任务完成事件: {task.task_id}")
+
+            async def task_completion_callback(completed_task: Task):
+                logger.info(f"任务完成回调被触发: {completed_task.task_id}, 状态: {completed_task.status}")
+                if completed_task.task_id == task.task_id and (
+                    completed_task.status == TaskStatus.COMPLETED or completed_task.status == TaskStatus.FAILED
+                ):
+                    logger.info(f"设置任务完成事件: {task.task_id}")
+                    task_completed.set()
+
+            # 添加回调
+            await task_manager.add_callback(task.task_id, task_completion_callback)
+            logger.info(f"已添加任务完成回调: {task.task_id}")
+
+            # 开始处理任务
+            await task_manager.process_task_sync(task.task_id)
+            logger.info(f"开始处理任务: {task.task_id}")
+
+            # 等待任务完成
+            try:
+                logger.info(f"等待任务完成: {task.task_id}")
+                await asyncio.wait_for(task_completed.wait(), timeout=1800)  # 30分钟超时
+                ##################sdsadsadsads
+                logger.info(f"任务完成事件已触发: {task.task_id}")
+            except asyncio.TimeoutError:
+                logger.error(f"任务处理超时: {task.task_id}")
+                raise HTTPException(status_code=500, detail="任务处理超时")
+
+            # 获取最终任务状态
+            result = await task_manager.get_task(task.task_id)
+            if not result:
+                raise HTTPException(status_code=404, detail="任务不存在")
+
+            if result.status == TaskStatus.COMPLETED:
+                # 获取输出文件路径
+                output_file = result.output_file_path
+                if not os.path.exists(output_file):
+                    logger.error(f"输出文件不存在: {output_file}")
+                    raise HTTPException(status_code=404, detail="输出文件不存在")
+
+                # 获取文件名
+                filename = os.path.basename(output_file)
+
+                # 创建文件流
+                async def file_iterator():
+                    async with aiofiles.open(output_file, "rb") as f:
+                        while chunk := await f.read(8192):  # 8KB chunks
+                            yield chunk
+
+                return StreamingResponse(
+                    file_iterator(), media_type="video/mp4", headers={"Content-Disposition": f'attachment; filename="{filename}"'}
+                )
+            else:
+                raise HTTPException(status_code=500, detail=result.error_message or "处理失败")
+
+        except Exception as e:
+            # 如果发生错误，立即清理临时文件
+            # if os.path.exists(temp_dir):
+            #     shutil.rmtree(temp_dir)
+            logger.error(f"创建任务失败: {str(e)}", exc_info=True)
+            raise HTTPException(status_code=500, detail=str(e))
+
+    except Exception as e:
+        logger.error(f"创建任务失败: {str(e)}", exc_info=True)
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.get("/tasks/{task_id}", response_model=Task)
+async def get_task(task_id: str):
+    """获取任务信息"""
+    task = await task_manager.get_task(task_id)
+    if not task:
+        raise HTTPException(status_code=404, detail="Task not found")
+    return task
+
+
+@app.get("/tasks/", response_model=List[Task])
+async def get_all_tasks():
+    """获取所有任务"""
+    return await task_manager.get_all_tasks()
+
+
+@app.post("/tasks/{task_id}/callback")
+async def add_task_callback(task_id: str, callback_url: str, background_tasks: BackgroundTasks):
+    """添加任务完成回调"""
+    if not await task_manager.get_task(task_id):
+        raise HTTPException(status_code=404, detail="Task not found")
+
+    async def callback_handler(task: Task):
+        try:
+            async with aiohttp.ClientSession() as session:
+                await session.post(callback_url, json=task.dict())
+        except Exception as e:
+            logger.error(f"Error sending callback to {callback_url}: {str(e)}")
+
+    await task_manager.add_callback(task_id, callback_handler)
+    return {"message": "Callback added successfully"}
+
+
+# 后台任务处理函数
+async def process_video_task(task_id: str, audio_file_path: str, video_file_path: str, out_file_path: str):
+    """处理视频任务的后台函数"""
+    try:
+        # 更新任务状态为处理中
+        await task_manager.update_task_status(task_id, TaskStatus.PROCESSING)
+
+        # 生成唯一的文件ID
+        file_id = f"{task_id}_{int(datetime.now().timestamp())}"
+
+        # 处理视频
+        success = await task_manager.process_manager.process_video(task_id, audio_file_path, video_file_path, out_file_path)
+
+        if success:
+            await task_manager.update_task_status(task_id, TaskStatus.COMPLETED)
+        else:
+            await task_manager.update_task_status(task_id, TaskStatus.FAILED, error_message="Video processing failed")
+    except Exception as e:
+        await task_manager.update_task_status(task_id, TaskStatus.FAILED, error_message=str(e))
+        # 清理临时文件
+        # if os.path.exists(temp_dir):
+        #     shutil.rmtree(temp_dir)
+        #     logger.info(f"已清理临时目录: {temp_dir}")
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run("api:app", host="0.0.0.0", port=8101, reload=True, workers=1, log_level="info")  # 开发模式下启用热重载  # 单进程模式