inpaintServer/app/api/endpoints.py

683 lines
24 KiB
Python

"""
API 엔드포인트
iopaint와 호환되는 인페인팅 및 배경 제거 API를 제공합니다.
"""
import time
import logging
from datetime import datetime
from typing import List
from fastapi import APIRouter, HTTPException, Response, Query, Request
from fastapi.responses import JSONResponse, StreamingResponse
import cv2
from ..core.config import settings
from ..core.worker_manager import worker_manager
from ..core.session_pool import session_pool
from ..models.schemas import (
InpaintRequest, RemoveBGRequest,
InpaintResponse, RemoveBGResponse, ServerConfigResponse,
HealthResponse, ModelInfo, Device, ResponseFormat, ImageFormat
)
from ..utils.image_utils import (
decode_base64_to_image, encode_image_to_base64, concat_alpha_channel,
pil_to_bytes, numpy_to_bytes, adjust_mask, gen_frontend_mask
)
import base64
import uuid
import io
from ..monitoring.dashboard import monitoring_data
from .stats import router as stats_router
from ..core.stats_manager import stats_manager
from ..core.batch_manager import batch_manager
from pathlib import Path
logger = logging.getLogger(__name__)
router = APIRouter()
router.include_router(stats_router, prefix="/api/v1", tags=["Stats"])
WEBHOOK_URL_FILE = Path(settings.PROJECT_ROOT) / "webhook_url.txt"
def get_webhook_url() -> str:
if WEBHOOK_URL_FILE.exists():
return WEBHOOK_URL_FILE.read_text().strip()
return ""
def save_webhook_url(url: str):
WEBHOOK_URL_FILE.write_text(url)
@router.get("/api/v1/webhook", summary="Get Discord Webhook URL")
async def get_webhook():
"""저장된 Discord 웹훅 URL을 조회합니다."""
return {"url": get_webhook_url()}
@router.post("/api/v1/webhook", summary="Set Discord Webhook URL")
async def set_webhook(request: Request):
"""Discord 웹훅 URL을 저장합니다."""
data = await request.json()
url = data.get("url", "")
save_webhook_url(url)
# Update settings in runtime if needed, or notifier can read from file directly
settings.DISCORD_WEBHOOK_URL = url
logger.info(f"Discord 웹훅 URL이 업데이트되었습니다: {url}")
return {"message": "Webhook URL이 성공적으로 저장되었습니다."}
def encode_image_to_format(image, format: ImageFormat = ImageFormat.png, quality: int = 95):
"""이미지를 지정된 형식으로 인코딩"""
if format == ImageFormat.png:
_, buffer = cv2.imencode('.png', image)
return buffer.tobytes(), 'image/png'
elif format == ImageFormat.webp:
_, buffer = cv2.imencode('.webp', image, [cv2.IMWRITE_WEBP_QUALITY, quality])
return buffer.tobytes(), 'image/webp'
elif format == ImageFormat.jpeg:
_, buffer = cv2.imencode('.jpg', image, [cv2.IMWRITE_JPEG_QUALITY, quality])
return buffer.tobytes(), 'image/jpeg'
else:
# 기본값은 PNG
_, buffer = cv2.imencode('.png', image)
return buffer.tobytes(), 'image/png'
def create_response(
image,
response_format: ResponseFormat = ResponseFormat.json,
image_format: ImageFormat = ImageFormat.png,
processing_time: float = 0.0,
success: bool = True,
error_msg: str = None
):
"""응답 형식에 따라 적절한 응답 생성"""
if not success:
if response_format == ResponseFormat.binary:
return Response(
content=b"",
status_code=500,
media_type="text/plain"
)
else:
return JSONResponse(
content={
"success": False,
"error": error_msg,
"processing_time": processing_time
},
status_code=500
)
if response_format == ResponseFormat.binary:
# 바이너리 응답 (기존 클라이언트 호환)
image_bytes, media_type = encode_image_to_format(image, image_format)
return Response(content=image_bytes, media_type=media_type)
elif response_format == ResponseFormat.stream:
# 스트리밍 응답
def generate():
image_bytes, _ = encode_image_to_format(image, image_format)
chunk_size = 8192
for i in range(0, len(image_bytes), chunk_size):
yield image_bytes[i:i + chunk_size]
_, media_type = encode_image_to_format(image, image_format)
return StreamingResponse(generate(), media_type=media_type)
elif response_format == ResponseFormat.base64:
# Base64 응답
image_b64 = encode_image_to_base64(image)
return JSONResponse(content={
"success": True,
"image": image_b64,
"processing_time": processing_time
})
else:
# JSON 응답 (기본값, iopaint 호환)
image_b64 = encode_image_to_base64(image)
return JSONResponse(content={
"success": True,
"image": image_b64,
"processing_time": processing_time
})
@router.get("/api/v1/health", response_model=HealthResponse, name="health_check")
async def health_check(request: Request):
"""서버 상태 확인"""
start_time = getattr(request.app.state, 'start_time', time.time())
current_time = time.time()
uptime_seconds = current_time - start_time
return HealthResponse(
status="ok",
uptime=uptime_seconds,
timestamp=datetime.fromtimestamp(current_time).isoformat(),
version="1.0.0"
)
# 호환용: 일부 클라이언트에서 /health 경로로 접근하는 경우가 있어 동일 응답 제공
@router.get("/health", response_model=HealthResponse, include_in_schema=False)
async def health_check_compat(request: Request):
return await health_check(request)
@router.get("/api/v1/realtime_status", include_in_schema=False)
async def get_realtime_status(request: Request):
"""실시간 세션/워커 상태 (모니터링 대시보드용)"""
try:
from app.core.session_pool import session_pool
from app.core.worker_manager import worker_manager
# 실시간 세션 상태 조회
session_status = session_pool.get_status()
# 실시간 워커 상태 조회
worker_status = worker_manager.get_status()
return {
"session_status": session_status,
"worker_status": worker_status,
"timestamp": time.time()
}
except Exception as e:
logger.error(f"실시간 상태 조회 실패: {e}")
return {
"session_status": {},
"worker_status": {},
"error": str(e),
"timestamp": time.time()
}
@router.get("/api/v1/server-config", response_model=ServerConfigResponse)
async def get_server_config():
"""서버 설정 정보 반환 (iopaint 호환)"""
try:
# 사용 가능한 모델 목록
models = [
ModelInfo(
name="simple-lama",
type="inpainting",
description="Simple LAMA 인페인팅 모델",
supported_formats=["png", "jpg", "jpeg"],
max_image_size=2048
),
ModelInfo(
name="migan",
type="inpainting",
description="MIGAN 인페인팅 모델",
supported_formats=["png", "jpg", "jpeg"],
max_image_size=2048
),
ModelInfo(
name="rembg",
type="rembg",
description="Rembg 배경 제거 모델",
supported_formats=["png", "jpg", "jpeg"],
max_image_size=2048
)
]
return ServerConfigResponse(
models=models,
max_file_size=settings.MAX_FILE_SIZE,
supported_formats=["png", "jpg", "jpeg"],
device=Device.cuda if settings.IS_JETSON else Device.cuda,
is_jetson=settings.IS_JETSON
)
except Exception as e:
logger.error(f"서버 설정 조회 실패: {e}")
raise HTTPException(status_code=500, detail=f"서버 설정 조회 실패: {str(e)}")
@router.post("/api/v1/inpaint")
async def inpaint_image(
request: InpaintRequest,
response_format: ResponseFormat = Query(ResponseFormat.binary, description="응답 형식 (기존 클라이언트 호환을 위해 기본값: binary)"),
image_format: ImageFormat = Query(ImageFormat.png, description="이미지 형식"),
http_request: Request = None,
):
"""인페인팅 API (iopaint 호환)"""
start_time = time.time()
alpha_channel = None # 변수 초기화
try:
req_id = f"req_{uuid.uuid4().hex[:8]}"
client_ip = None
try:
if http_request and http_request.client:
client_ip = http_request.client.host
except Exception:
client_ip = None
# 원본 base64 크기(바이트) 계산
try:
raw_img_b64 = request.image.split(',', 1)[1] if isinstance(request.image, str) and ',' in request.image else request.image
img_bytes_len = len(base64.b64decode(raw_img_b64)) if raw_img_b64 else 0
except Exception:
img_bytes_len = 0
try:
raw_mask_b64 = request.mask.split(',', 1)[1] if isinstance(request.mask, str) and ',' in request.mask else request.mask
mask_bytes_len = len(base64.b64decode(raw_mask_b64)) if raw_mask_b64 else 0
except Exception:
mask_bytes_len = 0
# base64 이미지 디코딩
image, alpha_channel, info, ext = decode_base64_to_image(request.image)
mask, _, mask_info, mask_ext = decode_base64_to_image(request.mask, gray=True)
# 이미지/마스크 메타 로깅
try:
img_h, img_w = image.shape[:2]
img_ch = image.shape[2] if len(image.shape) == 3 else 1
meta_image = {
"request_id": req_id,
"client_ip": client_ip,
"kind": "image",
"format": info.get("format"),
"mode": info.get("mode"),
"ext": ext,
"pil_size": info.get("size"),
"np_shape": tuple(image.shape),
"h": img_h,
"w": img_w,
"channels": img_ch,
"dtype": str(image.dtype),
"bytes": img_bytes_len,
"has_alpha": alpha_channel is not None,
}
logger.info(f"[INPAINT_META] {meta_image}")
except Exception:
pass
try:
mask_h, mask_w = mask.shape[:2]
mask_ch = mask.shape[2] if len(mask.shape) == 3 else 1
meta_mask = {
"request_id": req_id,
"client_ip": client_ip,
"kind": "mask",
"format": mask_info.get("format") if isinstance(mask_info, dict) else None,
"mode": mask_info.get("mode") if isinstance(mask_info, dict) else None,
"ext": mask_ext if 'mask_ext' in locals() else None,
"pil_size": mask_info.get("size") if isinstance(mask_info, dict) else None,
"np_shape": tuple(mask.shape),
"h": mask_h,
"w": mask_w,
"channels": mask_ch,
"dtype": str(mask.dtype),
"bytes": mask_bytes_len,
}
logger.info(f"[INPAINT_META] {meta_mask}")
except Exception:
pass
# alpha_channel이 None인 경우 기본값 설정
if alpha_channel is None:
alpha_channel = None # 투명도 채널이 없는 경우
# 이미지와 마스크 크기 검증
if image.shape[:2] != mask.shape[:2]:
raise HTTPException(
status_code=400,
detail=f"이미지 크기({image.shape[:2]})와 마스크 크기({mask.shape[:2]})가 일치하지 않습니다."
)
# 이미지 크기 검증
if not validate_image_size(image, settings.MAX_IMAGE_SIZE):
raise HTTPException(
status_code=400,
detail=f"이미지 크기가 너무 큽니다. 최대 {settings.MAX_IMAGE_SIZE}x{settings.MAX_IMAGE_SIZE}까지 지원합니다."
)
# 마스크 이진화
mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)[1]
# 모델 선택
model_name = request.model_name or "simple-lama"
# 워커에서 인페인팅 실행
if settings.USE_MICRO_BATCHING and model_name == "simple-lama":
# SimpleLama는 배치 관리자를 통해 처리
job_data = {
"image": image,
"mask": mask,
"prompt": request.prompt,
"negative_prompt": request.negative_prompt,
"sd_seed": request.sd_seed,
"num_inference_steps": request.num_inference_steps,
"guidance_scale": request.guidance_scale,
"strength": request.strength,
}
result = await batch_manager.add_job(job_data)
else:
# 다른 모델들은 기존 워커 매니저를 통해 단일 처리
result = await worker_manager.process_inpaint(
image=image,
mask=mask,
model_name=model_name,
prompt=request.prompt,
negative_prompt=request.negative_prompt,
sd_seed=request.sd_seed,
num_inference_steps=request.num_inference_steps,
guidance_scale=request.guidance_scale,
strength=request.strength
)
if result is None:
raise HTTPException(status_code=500, detail="인페인팅 처리 실패")
# 결과 이미지 처리
result_image = concat_alpha_channel(result, alpha_channel)
processing_time = time.time() - start_time
# 모니터링 통계 업데이트
monitoring_data.update_api_stats(
endpoint="/api/v1/inpaint",
success=True,
response_time=processing_time * 1000 # ms로 변환
)
# 응답 형식에 따라 적절한 응답 생성
return create_response(
image=result_image,
response_format=response_format,
image_format=image_format,
processing_time=processing_time,
success=True
)
except HTTPException:
# HTTPException은 상세 사유를 에러 로그에 남길 수 있도록 재전파됨 (미들웨어에서 잡혀 JSONL 기록)
raise
except Exception as e:
logger.error(f"인페인팅 처리 실패: {e}")
# 모니터링 통계 업데이트
processing_time = time.time() - start_time
monitoring_data.update_api_stats(
endpoint="/api/v1/inpaint",
success=False,
response_time=processing_time * 1000,
error=str(e)
)
# 응답 형식에 따라 적절한 에러 응답 생성
return create_response(
image=None,
response_format=response_format,
image_format=image_format,
processing_time=processing_time,
success=False,
error_msg=f"인페인팅 처리 실패: {str(e)}"
)
@router.post("/api/v1/remove_bg")
async def remove_background(
request: RemoveBGRequest,
response_format: ResponseFormat = Query(ResponseFormat.base64, description="응답 형식"),
image_format: ImageFormat = Query(ImageFormat.png, description="이미지 형식"),
http_request: Request = None,
):
"""배경 제거 API (iopaint 호환)"""
start_time = time.time()
alpha_channel = None # 변수 초기화
try:
req_id = f"req_{uuid.uuid4().hex[:8]}"
client_ip = None
try:
if http_request and http_request.client:
client_ip = http_request.client.host
except Exception:
client_ip = None
# 원본 base64 크기(바이트) 계산
try:
raw_img_b64 = request.image.split(',', 1)[1] if isinstance(request.image, str) and ',' in request.image else request.image
img_bytes_len = len(base64.b64decode(raw_img_b64)) if raw_img_b64 else 0
except Exception:
img_bytes_len = 0
# base64 이미지 디코딩
image, alpha_channel, info, ext = decode_base64_to_image(request.image)
# 이미지 메타 로깅
try:
img_h, img_w = image.shape[:2]
img_ch = image.shape[2] if len(image.shape) == 3 else 1
meta_image = {
"request_id": req_id,
"client_ip": client_ip,
"kind": "image",
"format": info.get("format"),
"mode": info.get("mode"),
"ext": ext,
"pil_size": info.get("size"),
"np_shape": tuple(image.shape),
"h": img_h,
"w": img_w,
"channels": img_ch,
"dtype": str(image.dtype),
"bytes": img_bytes_len,
"has_alpha": alpha_channel is not None,
}
logger.info(f"[REMOVEBG_META] {meta_image}")
except Exception:
pass
# alpha_channel이 None인 경우 기본값 설정
if alpha_channel is None:
alpha_channel = None # 투명도 채널이 없는 경우
# 이미지 크기 검증
if not validate_image_size(image, settings.MAX_IMAGE_SIZE):
raise HTTPException(
status_code=400,
detail=f"이미지 크기가 너무 큽니다. 최대 {settings.MAX_IMAGE_SIZE}x{settings.MAX_IMAGE_SIZE}까지 지원합니다."
)
# 모델 선택
model_name = request.model_name
if not model_name or model_name == "rembg":
model_name = "briaaiRMBG-1.4"
# 워커에서 배경 제거 실행
result_image, result_mask = await worker_manager.process_remove_bg(
image=image,
model_name=model_name
)
if result_image is None or result_mask is None:
raise HTTPException(status_code=500, detail="배경 제거 처리 실패")
processing_time = time.time() - start_time
# 모니터링 통계 업데이트
monitoring_data.update_api_stats(
endpoint="/api/v1/remove_bg",
success=True,
response_time=processing_time * 1000
)
# 응답 형식에 따라 적절한 응답 생성
if response_format == ResponseFormat.binary:
# 바이너리 응답은 이미지만 반환 (기존 클라이언트 호환)
return create_response(
image=result_image,
response_format=response_format,
image_format=image_format,
processing_time=processing_time,
success=True
)
elif response_format == ResponseFormat.base64:
# Base64 응답은 이미지만 반환
return create_response(
image=result_image,
response_format=response_format,
image_format=image_format,
processing_time=processing_time,
success=True
)
else:
# JSON 응답은 이미지와 마스크 모두 반환
result_base64 = encode_image_to_base64(result_image, ext)
mask_base64 = encode_image_to_base64(result_mask, "PNG")
return JSONResponse(content={
"success": True,
"image": result_base64,
"mask": mask_base64,
"processing_time": processing_time
})
except HTTPException:
# HTTPException은 상세 사유를 에러 로그에 남길 수 있도록 재전파됨
raise
except Exception as e:
logger.error(f"배경 제거 처리 실패: {e}")
# 모니터링 통계 업데이트
processing_time = time.time() - start_time
monitoring_data.update_api_stats(
endpoint="/api/v1/remove_bg",
success=False,
response_time=processing_time * 1000,
error=str(e)
)
# 응답 형식에 따라 적절한 에러 응답 생성
return create_response(
image=None,
response_format=response_format,
image_format=image_format,
processing_time=processing_time,
success=False,
error_msg=f"배경 제거 처리 실패: {str(e)}"
)
@router.post("/api/v1/run_plugin_gen_image", name="run_plugin_gen_image")
async def run_plugin_gen_image(
request: RemoveBGRequest,
response_format: ResponseFormat = Query(ResponseFormat.binary, description="응답 형식"),
image_format: ImageFormat = Query(ImageFormat.png, description="이미지 형식")
):
"""
IOPaint 플러그인 호환성을 위한 배경 제거 엔드포인트입니다.
내부적으로 /api/v1/remove_bg와 동일하게 동작합니다.
"""
logger.info(f"플러그인 호환 엔드포인트 '/api/v1/run_plugin_gen_image' 호출됨 (모델: {request.model_name})")
return await remove_background(request, response_format, image_format)
@router.get("/api/v1/model")
async def get_model_info():
"""모델 정보 반환 (클라이언트 헬스체크 호환)"""
start_time = time.time()
try:
# 현재 사용 가능한 모델 목록
models = [
{
"name": "simple-lama",
"type": "inpainting",
"status": "available",
"description": "Simple LAMA 인페인팅 모델"
},
{
"name": "migan",
"type": "inpainting",
"status": "available",
"description": "MIGAN 인페인팅 모델"
},
{
"name": "rembg",
"type": "rembg",
"status": "available",
"description": "Rembg 배경 제거 모델"
}
]
processing_time = time.time() - start_time
# 모니터링 통계 업데이트
monitoring_data.update_api_stats(
endpoint="/api/v1/model",
success=True,
response_time=processing_time * 1000 # ms로 변환
)
return {
"success": True,
"models": models,
"server_status": "running",
"version": "1.0.0",
"device": "cuda" if settings.IS_JETSON else "cuda",
"is_jetson": settings.IS_JETSON,
"timestamp": datetime.now().isoformat()
}
except Exception as e:
logger.error(f"모델 정보 조회 실패: {e}")
# 모니터링 통계 업데이트
processing_time = time.time() - start_time
monitoring_data.update_api_stats(
endpoint="/api/v1/model",
success=False,
response_time=processing_time * 1000,
error=str(e)
)
raise HTTPException(status_code=500, detail=f"모델 정보 조회 실패: {str(e)}")
@router.get("/api/v1/samplers")
async def get_samplers():
"""사용 가능한 샘플러 목록 반환 (iopaint 호환)"""
return [
"euler",
"euler_a",
"heun",
"dpm_2",
"dpm_2_a",
"lms",
"dpm_fast",
"dpm_adaptive",
"dpmpp_2s_a",
"dpmpp_sde",
"dpmpp_2m",
"ddim",
"uni_pc",
"uni_pc_bh2"
]
@router.get("/")
async def root():
"""루트 엔드포인트"""
return {
"message": "인페인팅 서버 API (iopaint 호환)",
"version": "1.0.0",
"docs": "/docs",
"health": "/api/v1/health"
}
# 유틸리티 함수
def validate_image_size(image, max_size):
"""이미지 크기 검증"""
try:
h, w = image.shape[:2]
pixels = h * w
max_pixels = max_size * max_size
return pixels <= max_pixels
except:
return False