49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
from fastapi import APIRouter, HTTPException
|
|
|
|
from app.schemas import GenerateRequest, HealthResponse, TaskResultResponse, TaskStatusResponse
|
|
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
@router.post("/generate", response_model=TaskStatusResponse)
|
|
async def create_generate_task(req: GenerateRequest):
|
|
task_manager = router.task_manager
|
|
return await task_manager.create_task(req)
|
|
|
|
|
|
@router.get("/tasks/{task_id}", response_model=TaskStatusResponse)
|
|
def get_task_status(task_id: str):
|
|
task_manager = router.task_manager
|
|
try:
|
|
return task_manager.get_status(task_id)
|
|
except KeyError as exc:
|
|
raise HTTPException(status_code=404, detail=str(exc)) from exc
|
|
|
|
|
|
@router.get("/tasks/{task_id}/result", response_model=TaskResultResponse)
|
|
def get_task_result(task_id: str):
|
|
task_manager = router.task_manager
|
|
try:
|
|
return task_manager.get_result(task_id)
|
|
except KeyError as exc:
|
|
raise HTTPException(status_code=404, detail=str(exc)) from exc
|
|
|
|
|
|
@router.get("/health", response_model=HealthResponse)
|
|
def health_check():
|
|
torch = router.torch
|
|
ltx_backend = router.ltx_backend
|
|
hunyuan_backend = router.hunyuan_backend
|
|
|
|
cuda_ok = bool(torch.cuda.is_available())
|
|
gpu_name = torch.cuda.get_device_name(0) if cuda_ok else None
|
|
|
|
return HealthResponse(
|
|
service_status="ok",
|
|
cuda_available=cuda_ok,
|
|
gpu_name=gpu_name,
|
|
ltx_loaded=ltx_backend.is_loaded(),
|
|
hunyuan_loaded=hunyuan_backend.is_loaded(),
|
|
)
|