inpaintServer/app/models/schemas.py

103 lines
4.6 KiB
Python

"""
API 스키마 정의
iopaint와 호환되는 입출력 형식을 지원합니다.
"""
from typing import Optional, List, Dict, Any, Tuple
from pydantic import BaseModel, Field
from enum import Enum
class Device(str, Enum):
cpu = "cpu"
cuda = "cuda"
mps = "mps"
class ResponseFormat(str, Enum):
"""응답 형식"""
json = "json" # JSON 형식 (기본값, iopaint 호환)
binary = "binary" # 바이너리 이미지 (기존 클라이언트 호환)
stream = "stream" # 스트리밍 응답
base64 = "base64" # Base64 문자열 (테스트 및 일부 클라이언트용)
class ImageFormat(str, Enum):
"""이미지 포맷"""
png = "png" # PNG (기본값, 무손실)
webp = "webp" # WebP (압축률 좋음)
jpeg = "jpeg" # JPEG (호환성 좋음)
class InpaintRequest(BaseModel):
"""인페인팅 요청 스키마 (iopaint 호환)"""
image: str = Field(..., description="base64로 인코딩된 원본 이미지")
mask: str = Field(..., description="base64로 인코딩된 마스크 이미지")
sd_seed: Optional[int] = Field(-1, description="Stable Diffusion 시드")
prompt: Optional[str] = Field("", description="프롬프트")
negative_prompt: Optional[str] = Field("", description="네거티브 프롬프트")
num_inference_steps: Optional[int] = Field(20, description="추론 스텝 수")
guidance_scale: Optional[float] = Field(7.5, description="가이던스 스케일")
strength: Optional[float] = Field(1.0, description="인페인팅 강도")
model_name: Optional[str] = Field("simple-lama", description="사용할 모델명")
model_config = {
"protected_namespaces": () # 보호 네임스페이스 비활성화
}
class RemoveBGRequest(BaseModel):
"""배경 제거 요청 스키마 (iopaint 호환)"""
image: str = Field(..., description="base64로 인코딩된 이미지")
model_name: Optional[str] = Field("rembg", description="사용할 모델명")
model_config = {
"protected_namespaces": () # 보호 네임스페이스 비활성화
}
class InpaintResponse(BaseModel):
"""인페인팅 응답 스키마 (iopaint 호환)"""
success: bool = Field(..., description="성공 여부")
image: Optional[str] = Field(None, description="base64로 인코딩된 결과 이미지")
error: Optional[str] = Field(None, description="에러 메시지")
processing_time: Optional[float] = Field(None, description="처리 시간 (초)")
seed: Optional[int] = Field(None, description="사용된 시드")
class RemoveBGResponse(BaseModel):
"""배경 제거 응답 스키마 (iopaint 호환)"""
success: bool = Field(..., description="성공 여부")
image: Optional[str] = Field(None, description="base64로 인코딩된 결과 이미지")
mask: Optional[str] = Field(None, description="base64로 인코딩된 마스크")
error: Optional[str] = Field(None, description="에러 메시지")
processing_time: Optional[float] = Field(None, description="처리 시간 (초)")
class ModelInfo(BaseModel):
"""모델 정보 스키마 (iopaint 호환)"""
name: str = Field(..., description="모델명")
type: str = Field(..., description="모델 타입 (inpainting/rembg)")
description: Optional[str] = Field(None, description="모델 설명")
supported_formats: List[str] = Field(default_factory=list, description="지원하는 이미지 형식")
max_image_size: Optional[int] = Field(None, description="최대 이미지 크기")
class ServerConfigResponse(BaseModel):
"""서버 설정 응답 스키마 (iopaint 호환)"""
models: List[ModelInfo] = Field(..., description="사용 가능한 모델 목록")
max_file_size: int = Field(..., description="최대 파일 크기 (MB)")
supported_formats: List[str] = Field(..., description="지원하는 이미지 형식")
device: Device = Field(..., description="현재 사용 중인 디바이스")
is_jetson: bool = Field(..., description="Jetson 시스템 여부")
class HealthResponse(BaseModel):
"""헬스 체크 응답 스키마 (iopaint 호환)"""
status: str = Field(..., description="서버 상태")
timestamp: str = Field(..., description="현재 시간")
version: str = Field(..., description="서버 버전")
uptime: float = Field(..., description="가동 시간 (초)")
class ErrorResponse(BaseModel):
"""에러 응답 스키마 (iopaint 호환)"""
error: str = Field(..., description="에러 메시지")
detail: Optional[str] = Field(None, description="상세 에러 정보")
timestamp: str = Field(..., description="에러 발생 시간")