'commit'

dsLightRag/Routes/JiMengRoute.py

@@ -3,11 +3,10 @@ import logging
 import fastapi
 from fastapi import APIRouter
 from fastapi import HTTPException
-from openai import AsyncOpenAI
 
-from Config import Config
-from JiMeng.Kit.JmTxt2ImgUtil import JmTxt2Img
 from JiMeng.Kit.FenJingTouGenerator import FenJingTouGenerator  # 导入分镜头生成器
+from JiMeng.Kit.JmImg2VideoUtil import JmImg2Video  # 导入视频生成工具
+from JiMeng.Kit.JmTxt2ImgUtil import JmTxt2Img
 
 # 创建路由路由器
 router = APIRouter(prefix="/api/jimeng", tags=["即梦"])
@@ -15,12 +14,6 @@ router = APIRouter(prefix="/api/jimeng", tags=["即梦"])
 # 配置日志
 logger = logging.getLogger(__name__)
 
-# 初始化异步 OpenAI 客户端
-client = AsyncOpenAI(
-    api_key=Config.ALY_LLM_API_KEY,
-    base_url=Config.ALY_LLM_BASE_URL
-)
-
 
 @router.post("/prompt_input")
 async def prompt_input(request: fastapi.Request):
@@ -87,3 +80,68 @@ async def generate_fenjingtou(request: fastapi.Request):
         logger.error(f"分镜头脚本生成失败: {str(e)}")
         raise HTTPException(status_code=500, detail=f"分镜头脚本生成失败: {str(e)}")
 
+
+@router.post("/create_video_task")
+async def create_video_task(request: fastapi.Request):
+    try:
+        data = await request.json()
+        image_url = data.get("image_url")
+        prompt = data.get("prompt")
+
+        if not image_url:
+            raise HTTPException(status_code=400, detail="缺少图片地址参数")
+        if not prompt:
+            raise HTTPException(status_code=400, detail="缺少分镜头脚本参数")
+
+        logger.info(f"收到视频任务创建请求，图片地址: {image_url}，分镜头脚本: {prompt}")
+
+        # 创建视频生成实例并调用方法
+        img2video = JmImg2Video()
+        task_id = img2video.create_video_task([image_url], prompt)
+
+        logger.info(f"视频任务创建成功，任务ID: {task_id}")
+
+        return {
+            "code": 200,
+            "message": "成功",
+            "data": {
+                "task_id": task_id
+            }
+        }
+    except HTTPException as e:
+        logger.error(f"请求参数错误: {str(e.detail)}")
+        raise e
+    except Exception as e:
+        logger.error(f"视频任务创建失败: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"视频任务创建失败: {str(e)}")
+
+
+@router.post("/query_video_task")
+async def query_video_task(request: fastapi.Request):
+    try:
+        data = await request.json()
+        task_id = data.get("task_id")
+
+        if not task_id:
+            raise HTTPException(status_code=400, detail="缺少任务ID参数")
+
+        logger.info(f"收到视频任务查询请求，任务ID: {task_id}")
+
+        # 创建视频生成实例并调用方法
+        img2video = JmImg2Video()
+        status_info = img2video.query_task_status(task_id)
+
+        logger.info(f"视频任务查询成功，状态: {status_info['status']}")
+
+        return {
+            "code": 200,
+            "message": "成功",
+            "data": status_info
+        }
+    except HTTPException as e:
+        logger.error(f"请求参数错误: {str(e.detail)}")
+        raise e
+    except Exception as e:
+        logger.error(f"视频任务查询失败: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"视频任务查询失败: {str(e)}")
+