43 lines
945 B
Python
43 lines
945 B
Python
from __future__ import annotations
|
|
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from PIL import Image
|
|
|
|
from ..image_io import DownloadedImage
|
|
|
|
|
|
class UpscaleEngineUnavailable(RuntimeError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class UpscaleResult:
|
|
image: Image.Image
|
|
metadata: dict[str, Any]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class EngineHealth:
|
|
status: str
|
|
engine: str
|
|
device: str
|
|
models_loaded: bool
|
|
details: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
class UpscaleEngine(ABC):
|
|
@abstractmethod
|
|
def health(self) -> EngineHealth:
|
|
raise NotImplementedError
|
|
|
|
def available(self) -> bool:
|
|
health = self.health()
|
|
|
|
return health.status == "ok" and health.models_loaded
|
|
|
|
@abstractmethod
|
|
def upscale(self, downloaded: DownloadedImage, scale: int, mode: str, output_format: str) -> UpscaleResult:
|
|
raise NotImplementedError |