Joey
71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
"""套图任务 API：轮询进度 / 导出 ZIP（进程内内存任务表）。"""
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
import zipfile
|
|
|
|
from fastapi import APIRouter, HTTPException
|
|
from fastapi.responses import StreamingResponse
|
|
|
|
from schemas import SuiteImageOut, SuiteOut
|
|
from services import storage
|
|
from services.tasks import IMG_OK, Task, get_task
|
|
|
|
router = APIRouter(prefix="/api", tags=["suites"])
|
|
|
|
|
|
def _task_out(task: Task) -> SuiteOut:
|
|
return SuiteOut(
|
|
id=task.id,
|
|
status=task.status,
|
|
style_set=task.style_set,
|
|
platform=task.platform,
|
|
lang=task.lang,
|
|
ratio=task.ratio,
|
|
provider=task.provider,
|
|
model=task.model,
|
|
total=task.total,
|
|
images=[
|
|
SuiteImageOut(
|
|
type_id=i.type_id, name=i.name, url=i.url,
|
|
status=i.status, error=i.error,
|
|
) for i in task.images
|
|
],
|
|
error=task.error,
|
|
)
|
|
|
|
|
|
@router.get("/suites/{suite_id}", response_model=SuiteOut)
|
|
async def get_suite(suite_id: str):
|
|
task = get_task(suite_id)
|
|
if task is None:
|
|
raise HTTPException(status_code=404, detail="任务不存在（服务可能已重启），请重新生成")
|
|
return _task_out(task)
|
|
|
|
|
|
@router.get("/suites/{suite_id}/zip")
|
|
async def download_suite_zip(suite_id: str):
|
|
"""把任务内所有成功图打包成 ZIP（中文文件名）。"""
|
|
task = get_task(suite_id)
|
|
if task is None:
|
|
raise HTTPException(status_code=404, detail="任务不存在（服务可能已重启）")
|
|
|
|
buf = io.BytesIO()
|
|
with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
|
|
seen: set[str] = set()
|
|
for i, img in enumerate([i for i in task.images if i.status == IMG_OK]):
|
|
path = storage.local_path(img.url or "")
|
|
if path is None:
|
|
continue
|
|
filename = img.name or img.type_id
|
|
if filename in seen: # 同类型多张时加序号防覆盖
|
|
filename = f"{filename}-{i + 1}"
|
|
seen.add(filename)
|
|
zf.write(path, f"{filename}{path.suffix or '.jpg'}")
|
|
buf.seek(0)
|
|
return StreamingResponse(
|
|
buf,
|
|
media_type="application/zip",
|
|
headers={"Content-Disposition": f"attachment; filename=\"suite-{suite_id}.zip\""},
|
|
)
|