58 lines
1.5 KiB
Python
58 lines
1.5 KiB
Python
|
"""Schemas for upscaling operations."""
|
||
|
|
from typing import Optional
|
||
|
|
|
||
|
|
from pydantic import BaseModel, Field
|
||
|
|
|
||
|
|
|
||
|
|
class UpscaleOptions(BaseModel):
|
||
|
|
"""Options for image upscaling."""
|
||
|
|
model: str = Field(
|
||
|
|
default='RealESRGAN_x4plus',
|
||
|
|
description='Model to use for upscaling (RealESRGAN_x2plus, RealESRGAN_x3plus, RealESRGAN_x4plus, etc.)'
|
||
|
|
)
|
||
|
|
tile_size: Optional[int] = Field(
|
||
|
|
default=None,
|
||
|
|
description='Tile size for processing large images to avoid OOM'
|
||
|
|
)
|
||
|
|
tile_pad: Optional[int] = Field(
|
||
|
|
default=None,
|
||
|
|
description='Padding between tiles'
|
||
|
|
)
|
||
|
|
outscale: Optional[float] = Field(
|
||
|
|
default=None,
|
||
|
|
description='Upsampling scale factor'
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class JobStatus(BaseModel):
|
||
|
|
"""Job status information."""
|
||
|
|
job_id: str
|
||
|
|
status: str # queued, processing, completed, failed
|
||
|
|
model: str
|
||
|
|
progress: float = Field(default=0.0, description='Progress as percentage 0-100')
|
||
|
|
result_path: Optional[str] = None
|
||
|
|
error: Optional[str] = None
|
||
|
|
created_at: str
|
||
|
|
started_at: Optional[str] = None
|
||
|
|
completed_at: Optional[str] = None
|
||
|
|
processing_time_seconds: Optional[float] = None
|
||
|
|
|
||
|
|
|
||
|
|
class UpscaleResult(BaseModel):
|
||
|
|
"""Upscaling result."""
|
||
|
|
success: bool
|
||
|
|
message: str
|
||
|
|
processing_time_seconds: float
|
||
|
|
model: str
|
||
|
|
input_size: Optional[tuple] = None
|
||
|
|
output_size: Optional[tuple] = None
|
||
|
|
|
||
|
|
|
||
|
|
class ModelInfo(BaseModel):
|
||
|
|
"""Information about an available model."""
|
||
|
|
name: str
|
||
|
|
scale: int
|
||
|
|
path: str
|
||
|
|
size_mb: float
|
||
|
|
available: bool
|