100 lines
3.8 KiB
Python
100 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
from contextlib import suppress
|
|
|
|
from fastapi import Depends, FastAPI, HTTPException, Request, status
|
|
from fastapi.responses import FileResponse, JSONResponse
|
|
|
|
from .config import Settings, get_settings
|
|
from .image_io import (
|
|
cleanup_expired_files,
|
|
delete_temp_file,
|
|
download_source_image,
|
|
ensure_directories,
|
|
resolve_result_path,
|
|
save_output_image,
|
|
)
|
|
from .schemas import HealthResponse, UpscaleRequest, UpscaleResponse
|
|
from .security import verify_bearer_token
|
|
from .upscaler import UpscaleEngineUnavailable, build_upscaler
|
|
|
|
|
|
def create_app(settings: Settings | None = None) -> FastAPI:
|
|
app = FastAPI(title="skinbase-enhance-worker", version="1.0.0")
|
|
resolved_settings = settings or get_settings()
|
|
ensure_directories(resolved_settings)
|
|
cleanup_expired_files(resolved_settings)
|
|
app.state.settings = resolved_settings
|
|
app.state.upscaler = build_upscaler(resolved_settings)
|
|
|
|
@app.get("/health", response_model=HealthResponse)
|
|
def health() -> HealthResponse:
|
|
engine_health = app.state.upscaler.health()
|
|
|
|
with suppress(Exception):
|
|
cleanup_expired_files(app.state.settings)
|
|
|
|
return HealthResponse(
|
|
status=engine_health.status,
|
|
service="skinbase-enhance-worker",
|
|
engine=engine_health.engine,
|
|
device=engine_health.device,
|
|
models_loaded=engine_health.models_loaded,
|
|
max_input_width=app.state.settings.max_input_width,
|
|
max_input_height=app.state.settings.max_input_height,
|
|
max_output_width=app.state.settings.max_output_width,
|
|
max_output_height=app.state.settings.max_output_height,
|
|
realesrgan=engine_health.details.get("realesrgan"),
|
|
)
|
|
|
|
@app.post("/v1/upscale", response_model=UpscaleResponse)
|
|
def upscale(payload: UpscaleRequest, request: Request, _: None = Depends(verify_bearer_token)):
|
|
cleanup_expired_files(app.state.settings)
|
|
downloaded = None
|
|
|
|
try:
|
|
downloaded = download_source_image(payload.source_url, app.state.settings)
|
|
result = app.state.upscaler.upscale(downloaded, payload.scale, payload.mode, payload.output_format)
|
|
stored = save_output_image(result.image, payload.output_format, app.state.settings, payload.job_id)
|
|
|
|
return UpscaleResponse(
|
|
success=True,
|
|
job_id=payload.job_id,
|
|
output_url=str(request.base_url).rstrip("/") + f"/v1/results/{stored.filename}",
|
|
width=stored.width,
|
|
height=stored.height,
|
|
filesize=stored.filesize,
|
|
mime=stored.mime,
|
|
metadata=result.metadata,
|
|
)
|
|
except HTTPException as exc:
|
|
return JSONResponse(status_code=exc.status_code, content={"success": False, "error": exc.detail})
|
|
except UpscaleEngineUnavailable as exc:
|
|
return JSONResponse(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, content={"success": False, "error": str(exc)})
|
|
finally:
|
|
delete_temp_file(downloaded.path if downloaded is not None else None)
|
|
|
|
@app.get("/v1/results/{filename}")
|
|
def result(filename: str):
|
|
path = resolve_result_path(app.state.settings, filename)
|
|
|
|
if not path.exists() or not path.is_file():
|
|
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Not found")
|
|
|
|
return FileResponse(path)
|
|
|
|
@app.delete("/v1/results/{filename}")
|
|
def delete_result(filename: str, _: None = Depends(verify_bearer_token)):
|
|
path = resolve_result_path(app.state.settings, filename)
|
|
|
|
if not path.exists() or not path.is_file():
|
|
return {"success": True, "deleted": False}
|
|
|
|
path.unlink(missing_ok=True)
|
|
|
|
return {"success": True, "deleted": True}
|
|
|
|
return app
|
|
|
|
|
|
app = create_app() |