dsProject/dsLightRag/Routes/JiMengRoute.py

import logging

import fastapi
from fastapi import APIRouter
from fastapi import HTTPException

from JiMeng.Kit.FenJingTouGenerator import FenJingTouGenerator  # 导入分镜头生成器
from JiMeng.Kit.JmImg2VideoUtil import JmImg2Video  # 导入视频生成工具
from JiMeng.Kit.JmTxt2ImgUtil import JmTxt2Img
import os
import uuid
import requests
from Util.ObsUtil import ObsUploader

# 创建路由路由器
router = APIRouter(prefix="/api/jimeng", tags=["即梦"])

# 配置日志
logger = logging.getLogger(__name__)


@router.post("/prompt_input")
async def prompt_input(request: fastapi.Request):
    try:
        data = await request.json()
        prompt = data.get("prompt")
        if not prompt:
            raise HTTPException(status_code=400, detail="缺少提示词参数")

        logger.info(f"收到图片生成请求，提示词: {prompt}")
        # 调用 JmTxt2Img 生成图片
        image_url = JmTxt2Img.generate_image(prompt)
        logger.info(f"图片生成成功，URL: {image_url}")

        return {
            "code": 200,
            "message": "成功",
            "data": {
                "image_url": image_url
            }
        }
    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("/generate_fenjingtou")
async def generate_fenjingtou(request: fastapi.Request):
    try:
        data = await request.json()
        image_url = data.get("image_url")
        prompt_text = data.get("prompt_text")

        if not image_url:
            raise HTTPException(status_code=400, detail="缺少图片地址参数")
        if not prompt_text:
            raise HTTPException(status_code=400, detail="缺少提示词参数")

        logger.info(f"收到分镜头生成请求，图片地址: {image_url}，提示词: {prompt_text}")

        # 创建分镜头生成器实例并调用方法
        generator = FenJingTouGenerator()
        fenjingtou_content = generator.generate_fen_jing_tou(image_url, prompt_text)

        if not fenjingtou_content:
            raise HTTPException(status_code=500, detail="分镜头脚本生成失败")

        logger.info(f"分镜头脚本生成成功")

        return {
            "code": 200,
            "message": "成功",
            "data": {
                "fenjingtou_content": fenjingtou_content
            }
        }
    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("/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)}")


@router.post("/upload_url_to_obs")
async def upload_url_to_obs(request: fastapi.Request):
    try:
        data = await request.json()
        url = data.get("url")

        if not url:
            raise HTTPException(status_code=400, detail="缺少URL参数")

        logger.info(f"收到URL上传请求，URL: {url}")

        # 使用UUID生成唯一文件名
        unique_id = uuid.uuid4()
        object_key = f"HuangHai/JiMeng/{unique_id}.mp4"

        # 临时文件保存路径
        temp_file_path = os.path.join(os.path.dirname(__file__), f"{unique_id}.mp4")

        try:
            # 下载URL内容到临时文件
            logger.info(f"开始下载URL内容: {url}")
            with requests.get(url, stream=True, timeout=3600) as r:
                r.raise_for_status()
                with open(temp_file_path, 'wb') as f:
                    for chunk in r.iter_content(chunk_size=8192):
                        f.write(chunk)
            logger.info(f"URL内容下载完成，保存至: {temp_file_path}")

            # 上传文件到OBS
            obs_uploader = ObsUploader()
            logger.info(f"开始上传文件到OBS: {temp_file_path}\n对象键: {object_key}")

            # 正确处理元组返回值
            success, response_info = obs_uploader.upload_file(
                file_path=temp_file_path,
                object_key=object_key
            )

            if success:
                # 构造OBS URL
                from Config.Config import OBS_SERVER, OBS_BUCKET
                obs_url = f"https://{OBS_BUCKET}.{OBS_SERVER}/{object_key}"
                logger.info(f"文件上传成功，OBS URL: {obs_url}")
                return {
                    "code": 200,
                    "message": "成功",
                    "data": {
                        "obs_url": obs_url,
                        "object_key": object_key
                    }
                }
            else:
                error_msg = f"上传失败: {str(response_info)}"
                logger.error(error_msg)
                raise HTTPException(status_code=500, detail=error_msg)

        finally:
            # 清理临时文件
            if os.path.exists(temp_file_path):
                os.remove(temp_file_path)
                logger.info(f"临时文件已删除: {temp_file_path}")

    except HTTPException as e:
        logger.error(f"请求参数错误: {str(e.detail)}")
        raise e
    except Exception as e:
        logger.error(f"URL上传处理失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"URL上传处理失败: {str(e)}")