112 lines
2.6 KiB
Python
112 lines
2.6 KiB
Python
|
"""Real-ESRGAN API application."""
|
||
|
|
import logging
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
from contextlib import asynccontextmanager
|
||
|
|
|
||
|
|
from fastapi import FastAPI
|
||
|
|
from fastapi.middleware.cors import CORSMiddleware
|
||
|
|
|
||
|
|
# Ensure app is importable
|
||
|
|
_app_path = os.path.dirname(__file__)
|
||
|
|
if _app_path not in sys.path:
|
||
|
|
sys.path.insert(0, _app_path)
|
||
|
|
|
||
|
|
from app.routers import health, models, upscale
|
||
|
|
from app.services import file_manager, realesrgan_bridge, worker
|
||
|
|
|
||
|
|
logging.basicConfig(
|
||
|
|
level=logging.INFO,
|
||
|
|
format='%(asctime)s %(levelname)s %(name)s: %(message)s'
|
||
|
|
)
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
def _process_upscale_job(job) -> None:
|
||
|
|
"""Worker function to process upscaling jobs."""
|
||
|
|
from app.services import realesrgan_bridge
|
||
|
|
|
||
|
|
bridge = realesrgan_bridge.get_bridge()
|
||
|
|
success, message, _ = bridge.upscale(
|
||
|
|
input_path=job.input_path,
|
||
|
|
output_path=job.output_path,
|
||
|
|
model_name=job.model,
|
||
|
|
outscale=job.outscale,
|
||
|
|
)
|
||
|
|
|
||
|
|
if not success:
|
||
|
|
raise Exception(message)
|
||
|
|
|
||
|
|
|
||
|
|
@asynccontextmanager
|
||
|
|
async def lifespan(app: FastAPI):
|
||
|
|
"""Application lifecycle manager."""
|
||
|
|
# Startup
|
||
|
|
logger.info('Starting Real-ESRGAN API...')
|
||
|
|
file_manager.ensure_directories()
|
||
|
|
|
||
|
|
bridge = realesrgan_bridge.get_bridge()
|
||
|
|
if not bridge.initialize():
|
||
|
|
logger.warning('Real-ESRGAN initialization failed (will attempt on first use)')
|
||
|
|
|
||
|
|
wq = worker.get_worker_queue(_process_upscale_job, num_workers=2)
|
||
|
|
wq.start()
|
||
|
|
|
||
|
|
logger.info('Real-ESRGAN API ready')
|
||
|
|
yield
|
||
|
|
|
||
|
|
# Shutdown
|
||
|
|
logger.info('Shutting down Real-ESRGAN API...')
|
||
|
|
wq.stop()
|
||
|
|
logger.info('Real-ESRGAN API stopped')
|
||
|
|
|
||
|
|
|
||
|
|
app = FastAPI(
|
||
|
|
title='Real-ESRGAN API',
|
||
|
|
version='1.0.0',
|
||
|
|
description='REST API for Real-ESRGAN image upscaling with async job processing',
|
||
|
|
lifespan=lifespan,
|
||
|
|
)
|
||
|
|
|
||
|
|
# CORS middleware
|
||
|
|
app.add_middleware(
|
||
|
|
CORSMiddleware,
|
||
|
|
allow_origins=['*'],
|
||
|
|
allow_credentials=True,
|
||
|
|
allow_methods=['*'],
|
||
|
|
allow_headers=['*'],
|
||
|
|
)
|
||
|
|
|
||
|
|
# Include routers
|
||
|
|
app.include_router(health.router)
|
||
|
|
app.include_router(models.router)
|
||
|
|
app.include_router(upscale.router)
|
||
|
|
|
||
|
|
|
||
|
|
@app.get('/')
|
||
|
|
async def root():
|
||
|
|
"""API root endpoint."""
|
||
|
|
return {
|
||
|
|
'name': 'Real-ESRGAN API',
|
||
|
|
'version': '1.0.0',
|
||
|
|
'docs': '/docs',
|
||
|
|
'redoc': '/redoc',
|
||
|
|
'endpoints': {
|
||
|
|
'health': '/api/v1/health',
|
||
|
|
'system': '/api/v1/system',
|
||
|
|
'models': '/api/v1/models',
|
||
|
|
'upscale': '/api/v1/upscale',
|
||
|
|
'jobs': '/api/v1/jobs',
|
||
|
|
},
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
import uvicorn
|
||
|
|
uvicorn.run(
|
||
|
|
'app.main:app',
|
||
|
|
host='0.0.0.0',
|
||
|
|
port=8000,
|
||
|
|
reload=False,
|
||
|
|
)
|