69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
from fastapi import APIRouter, HTTPException, Request
|
||
from fastapi.responses import StreamingResponse
|
||
from service import MidjourneyService
|
||
from utils import jingrow_api_verify_and_billing
|
||
from settings import settings
|
||
import json
|
||
import asyncio
|
||
from typing import AsyncGenerator, List
|
||
|
||
router = APIRouter(prefix=settings.router_prefix)
|
||
service = MidjourneyService()
|
||
|
||
@router.post(settings.generate_route)
|
||
@jingrow_api_verify_and_billing(api_name=settings.api_name)
|
||
async def generate_image(data: dict, request: Request):
|
||
"""
|
||
根据文本提示生成图像
|
||
|
||
Args:
|
||
data: 包含文本提示和配置参数的字典
|
||
request: FastAPI 请求对象
|
||
|
||
Returns:
|
||
生成的图像内容
|
||
"""
|
||
if "prompt" not in data:
|
||
raise HTTPException(status_code=400, detail="缺少prompt参数")
|
||
|
||
prompt = data["prompt"]
|
||
config = data.get("config", {})
|
||
|
||
async def generate() -> AsyncGenerator[str, None]:
|
||
async for result in service.generate_image(prompt, config):
|
||
yield json.dumps(result, ensure_ascii=False) + "\n"
|
||
|
||
return StreamingResponse(
|
||
generate(),
|
||
media_type="application/x-ndjson",
|
||
headers={"X-Content-Type-Options": "nosniff"}
|
||
)
|
||
|
||
@router.post(settings.batch_route)
|
||
@jingrow_api_verify_and_billing(api_name=settings.api_name)
|
||
async def batch_process_images(data: dict, request: Request):
|
||
"""
|
||
批量处理多个图像URL，将每张图片分割成4张并保存
|
||
|
||
Args:
|
||
data: 包含图片URLs列表的字典
|
||
request: FastAPI 请求对象
|
||
|
||
Returns:
|
||
处理结果的流式响应
|
||
"""
|
||
if "image_urls" not in data or not isinstance(data["image_urls"], list):
|
||
raise HTTPException(status_code=400, detail="缺少有效的image_urls参数")
|
||
|
||
image_urls: List[str] = data["image_urls"]
|
||
config = data.get("config", {})
|
||
|
||
async def process() -> AsyncGenerator[str, None]:
|
||
async for result in service.process_batch(image_urls, config):
|
||
yield json.dumps(result, ensure_ascii=False) + "\n"
|
||
|
||
return StreamingResponse(
|
||
process(),
|
||
media_type="application/x-ndjson",
|
||
headers={"X-Content-Type-Options": "nosniff"}
|
||
) |