107 lines
3.2 KiB
Python
107 lines
3.2 KiB
Python
import logging
|
||
import json
|
||
import time
|
||
import datetime
|
||
import requests
|
||
import uuid
|
||
import os
|
||
from typing import Optional, Dict, Any
|
||
from fastapi import APIRouter, HTTPException, Query, BackgroundTasks
|
||
from pydantic import BaseModel
|
||
from Midjourney.Txt2Img import Txt2Img
|
||
from Config import Config
|
||
import asyncio
|
||
import threading
|
||
from fastapi.responses import StreamingResponse
|
||
from fastapi import APIRouter, Request, HTTPException, BackgroundTasks
|
||
|
||
# 创建路由路由器
|
||
router = APIRouter(prefix="/api/mj", tags=["文生图"])
|
||
|
||
# 配置日志
|
||
logger = logging.getLogger(__name__)
|
||
|
||
class ImagineRequest(BaseModel):
|
||
prompt: str
|
||
base64_array: Optional[list] = None
|
||
notify_hook: Optional[str] = None
|
||
|
||
class TaskStatusResponse(BaseModel):
|
||
task_id: str
|
||
status: str
|
||
image_url: Optional[str] = None
|
||
progress: Optional[int] = None
|
||
error: Optional[str] = None
|
||
|
||
@router.post("/imagine", response_model=Dict[str, str])
|
||
async def submit_imagine(request: ImagineRequest):
|
||
"""
|
||
提交文生图请求
|
||
|
||
Args:
|
||
request: 包含提示词、垫图和回调地址的请求体
|
||
|
||
Returns:
|
||
包含MJ服务器任务ID的字典
|
||
"""
|
||
try:
|
||
# 直接提交到Midjourney并返回MJ任务ID
|
||
midjourney_task_id = Txt2Img.submit_imagine(
|
||
prompt=request.prompt,
|
||
base64_array=request.base64_array,
|
||
notify_hook=request.notify_hook
|
||
)
|
||
|
||
logger.info(f"提交文生图请求成功，MJ任务ID: {midjourney_task_id}, 提示词: {request.prompt}")
|
||
return {"task_id": midjourney_task_id}
|
||
except Exception as e:
|
||
logger.error(f"提交文生图请求失败: {str(e)}")
|
||
raise HTTPException(status_code=500, detail=f"提交文生图请求失败: {str(e)}")
|
||
|
||
@router.get("/task_status")
|
||
async def get_task_status(task_id: str):
|
||
"""
|
||
直接查询MJ服务器获取任务状态
|
||
|
||
Args:
|
||
task_id: MJ服务器返回的任务ID
|
||
|
||
Returns:
|
||
包含任务状态的字典
|
||
"""
|
||
try:
|
||
# 直接查询MJ服务器获取实时状态
|
||
result = Txt2Img.query_task_status(task_id)
|
||
|
||
# 处理查询结果
|
||
if result.get("status") == "SUCCESS":
|
||
return {
|
||
"task_id": task_id,
|
||
"status": "completed",
|
||
"image_url": result.get("imageUrl"),
|
||
"progress": 100,
|
||
"error": None
|
||
}
|
||
elif result.get("status") == "FAILED":
|
||
return {
|
||
"task_id": task_id,
|
||
"status": "failed",
|
||
"image_url": None,
|
||
"progress": 0,
|
||
"error": result.get("errorMsg", "未知错误")
|
||
}
|
||
else:
|
||
return {
|
||
"task_id": task_id,
|
||
"status": "processing",
|
||
"progress": result.get("progress", 0),
|
||
"error": None
|
||
}
|
||
except Exception as e:
|
||
logger.error(f"查询任务状态失败: {str(e)}")
|
||
return {
|
||
"task_id": task_id,
|
||
"status": "error",
|
||
"message": f"查询任务状态失败: {str(e)}"
|
||
}, 500
|