2023-01-28 20:56:06 +00:00
|
|
|
from enum import IntEnum
|
2023-01-28 04:48:06 +00:00
|
|
|
from typing import Any, Dict, Literal, Optional, Tuple, Union
|
|
|
|
|
|
|
|
|
2023-01-28 20:56:06 +00:00
|
|
|
class SizeChart(IntEnum):
|
2023-02-05 13:53:26 +00:00
|
|
|
mini = 128 # small tile for very expensive models
|
|
|
|
half = 256 # half tile for outpainting
|
|
|
|
auto = 512 # auto tile size
|
2023-01-28 20:56:06 +00:00
|
|
|
hd1k = 2**10
|
|
|
|
hd2k = 2**11
|
|
|
|
hd4k = 2**12
|
|
|
|
hd8k = 2**13
|
|
|
|
hd16k = 2**14
|
|
|
|
hd64k = 2**16
|
|
|
|
|
|
|
|
|
2023-02-12 00:00:18 +00:00
|
|
|
class TileOrder:
|
|
|
|
grid = "grid"
|
|
|
|
kernel = "kernel"
|
|
|
|
spiral = "spiral"
|
|
|
|
|
|
|
|
|
2023-01-28 04:48:06 +00:00
|
|
|
Param = Union[str, int, float]
|
|
|
|
Point = Tuple[int, int]
|
|
|
|
|
|
|
|
|
|
|
|
class Border:
|
|
|
|
def __init__(self, left: int, right: int, top: int, bottom: int) -> None:
|
|
|
|
self.left = left
|
|
|
|
self.right = right
|
|
|
|
self.top = top
|
|
|
|
self.bottom = bottom
|
|
|
|
|
2023-01-28 23:09:19 +00:00
|
|
|
def __str__(self) -> str:
|
2023-02-05 13:53:26 +00:00
|
|
|
return "%s %s %s %s" % (self.left, self.top, self.right, self.bottom)
|
2023-01-28 23:09:19 +00:00
|
|
|
|
2023-02-02 04:31:01 +00:00
|
|
|
def tojson(self):
|
|
|
|
return {
|
2023-02-05 13:53:26 +00:00
|
|
|
"left": self.left,
|
|
|
|
"right": self.right,
|
|
|
|
"top": self.top,
|
|
|
|
"bottom": self.bottom,
|
2023-02-02 04:31:01 +00:00
|
|
|
}
|
|
|
|
|
2023-01-28 15:55:47 +00:00
|
|
|
@classmethod
|
|
|
|
def even(cls, all: int):
|
|
|
|
return Border(all, all, all, all)
|
|
|
|
|
2023-01-28 04:48:06 +00:00
|
|
|
|
|
|
|
class Size:
|
|
|
|
def __init__(self, width: int, height: int) -> None:
|
|
|
|
self.width = width
|
|
|
|
self.height = height
|
|
|
|
|
2023-01-28 23:09:19 +00:00
|
|
|
def __str__(self) -> str:
|
2023-02-05 13:53:26 +00:00
|
|
|
return "%sx%s" % (self.width, self.height)
|
2023-01-28 23:09:19 +00:00
|
|
|
|
2023-01-28 04:48:06 +00:00
|
|
|
def add_border(self, border: Border):
|
2023-02-05 13:53:26 +00:00
|
|
|
return Size(
|
|
|
|
border.left + self.width + border.right,
|
|
|
|
border.top + self.height + border.right,
|
|
|
|
)
|
2023-01-28 04:48:06 +00:00
|
|
|
|
|
|
|
def tojson(self) -> Dict[str, int]:
|
|
|
|
return {
|
2023-02-05 13:53:26 +00:00
|
|
|
"height": self.height,
|
|
|
|
"width": self.width,
|
2023-01-28 04:48:06 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
|
2023-02-04 19:49:34 +00:00
|
|
|
class DeviceParams:
|
2023-02-05 13:53:26 +00:00
|
|
|
def __init__(
|
|
|
|
self, device: str, provider: str, options: Optional[dict] = None
|
|
|
|
) -> None:
|
2023-02-04 19:49:34 +00:00
|
|
|
self.device = device
|
|
|
|
self.provider = provider
|
|
|
|
self.options = options
|
|
|
|
|
2023-02-04 19:56:17 +00:00
|
|
|
def __str__(self) -> str:
|
2023-02-05 13:53:26 +00:00
|
|
|
return "%s - %s (%s)" % (self.device, self.provider, self.options)
|
2023-02-04 19:56:17 +00:00
|
|
|
|
2023-02-04 19:49:34 +00:00
|
|
|
def torch_device(self) -> str:
|
2023-02-05 13:53:26 +00:00
|
|
|
if self.device.startswith("cuda"):
|
2023-02-04 19:49:34 +00:00
|
|
|
return self.device
|
|
|
|
else:
|
2023-02-05 13:53:26 +00:00
|
|
|
return "cpu"
|
2023-02-04 19:49:34 +00:00
|
|
|
|
|
|
|
|
2023-01-28 04:48:06 +00:00
|
|
|
class ImageParams:
|
|
|
|
def __init__(
|
|
|
|
self,
|
|
|
|
model: str,
|
|
|
|
scheduler: Any,
|
|
|
|
prompt: str,
|
|
|
|
cfg: float,
|
|
|
|
steps: int,
|
2023-02-05 13:53:26 +00:00
|
|
|
seed: int,
|
2023-02-05 23:15:37 +00:00
|
|
|
negative_prompt: Optional[str] = None,
|
|
|
|
lpw: Optional[bool] = False,
|
2023-01-28 04:48:06 +00:00
|
|
|
) -> None:
|
|
|
|
self.model = model
|
|
|
|
self.scheduler = scheduler
|
|
|
|
self.prompt = prompt
|
|
|
|
self.negative_prompt = negative_prompt
|
|
|
|
self.cfg = cfg
|
|
|
|
self.seed = seed
|
2023-02-05 16:50:15 +00:00
|
|
|
self.steps = steps
|
2023-02-05 23:15:37 +00:00
|
|
|
self.lpw = lpw or False
|
2023-01-28 04:48:06 +00:00
|
|
|
|
2023-02-05 03:17:39 +00:00
|
|
|
def tojson(self) -> Dict[str, Optional[Param]]:
|
2023-01-28 04:48:06 +00:00
|
|
|
return {
|
2023-02-05 13:53:26 +00:00
|
|
|
"model": self.model,
|
|
|
|
"scheduler": self.scheduler.__name__,
|
|
|
|
"prompt": self.prompt,
|
|
|
|
"negativePrompt": self.negative_prompt,
|
2023-02-05 16:50:15 +00:00
|
|
|
"cfg": self.cfg,
|
|
|
|
"seed": self.seed,
|
2023-02-05 13:53:26 +00:00
|
|
|
"steps": self.steps,
|
2023-02-05 23:15:37 +00:00
|
|
|
"lpw": self.lpw,
|
2023-01-28 04:48:06 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
class StageParams:
|
2023-02-05 13:53:26 +00:00
|
|
|
"""
|
2023-01-28 04:48:06 +00:00
|
|
|
Parameters for a chained pipeline stage
|
2023-02-05 13:53:26 +00:00
|
|
|
"""
|
2023-01-28 04:48:06 +00:00
|
|
|
|
|
|
|
def __init__(
|
|
|
|
self,
|
|
|
|
name: Optional[str] = None,
|
|
|
|
outscale: int = 1,
|
2023-02-12 00:00:18 +00:00
|
|
|
tile_order: str = TileOrder.grid,
|
|
|
|
tile_size: int = SizeChart.auto,
|
2023-01-28 04:48:06 +00:00
|
|
|
# batch_size: int = 1,
|
|
|
|
) -> None:
|
|
|
|
self.name = name
|
|
|
|
self.outscale = outscale
|
2023-02-12 00:00:18 +00:00
|
|
|
self.tile_order = tile_order
|
|
|
|
self.tile_size = tile_size
|
2023-01-28 04:48:06 +00:00
|
|
|
|
|
|
|
|
2023-02-05 13:53:26 +00:00
|
|
|
class UpscaleParams:
|
2023-01-28 04:48:06 +00:00
|
|
|
def __init__(
|
|
|
|
self,
|
|
|
|
upscale_model: str,
|
|
|
|
correction_model: Optional[str] = None,
|
|
|
|
denoise: float = 0.5,
|
|
|
|
faces=True,
|
2023-02-06 23:13:37 +00:00
|
|
|
face_outscale: int = 1,
|
2023-01-28 04:48:06 +00:00
|
|
|
face_strength: float = 0.5,
|
2023-02-05 13:53:26 +00:00
|
|
|
format: Literal["onnx", "pth"] = "onnx",
|
2023-01-28 04:48:06 +00:00
|
|
|
half=False,
|
|
|
|
outscale: int = 1,
|
|
|
|
scale: int = 4,
|
|
|
|
pre_pad: int = 0,
|
|
|
|
tile_pad: int = 10,
|
|
|
|
) -> None:
|
|
|
|
self.upscale_model = upscale_model
|
|
|
|
self.correction_model = correction_model
|
|
|
|
self.denoise = denoise
|
|
|
|
self.faces = faces
|
2023-02-06 23:13:37 +00:00
|
|
|
self.face_outscale = face_outscale
|
2023-01-28 04:48:06 +00:00
|
|
|
self.face_strength = face_strength
|
|
|
|
self.format = format
|
|
|
|
self.half = half
|
|
|
|
self.outscale = outscale
|
|
|
|
self.pre_pad = pre_pad
|
|
|
|
self.scale = scale
|
|
|
|
self.tile_pad = tile_pad
|
|
|
|
|
2023-01-31 23:08:30 +00:00
|
|
|
def rescale(self, scale: int):
|
|
|
|
return UpscaleParams(
|
|
|
|
self.upscale_model,
|
|
|
|
correction_model=self.correction_model,
|
|
|
|
denoise=self.denoise,
|
|
|
|
faces=self.faces,
|
2023-02-06 23:13:37 +00:00
|
|
|
face_outscale=self.face_outscale,
|
2023-01-31 23:08:30 +00:00
|
|
|
face_strength=self.face_strength,
|
|
|
|
format=self.format,
|
|
|
|
half=self.half,
|
|
|
|
outscale=scale,
|
|
|
|
scale=scale,
|
|
|
|
pre_pad=self.pre_pad,
|
|
|
|
tile_pad=self.tile_pad,
|
|
|
|
)
|
|
|
|
|
2023-01-28 04:48:06 +00:00
|
|
|
def resize(self, size: Size) -> Size:
|
2023-02-06 23:59:34 +00:00
|
|
|
return Size(
|
|
|
|
size.width * self.outscale * self.face_outscale,
|
|
|
|
size.height * self.outscale * self.face_outscale,
|
|
|
|
)
|
2023-02-02 04:31:01 +00:00
|
|
|
|
|
|
|
def tojson(self):
|
|
|
|
return {
|
2023-02-05 16:50:15 +00:00
|
|
|
"upscale_model": self.upscale_model,
|
|
|
|
"correction_model": self.correction_model,
|
|
|
|
"denoise": self.denoise,
|
|
|
|
"faces": self.faces,
|
2023-02-06 23:13:37 +00:00
|
|
|
"face_outscale": self.face_outscale,
|
2023-02-05 16:50:15 +00:00
|
|
|
"face_strength": self.face_strength,
|
|
|
|
"format": self.format,
|
|
|
|
"half": self.half,
|
2023-02-05 13:53:26 +00:00
|
|
|
"outscale": self.outscale,
|
2023-02-05 16:50:15 +00:00
|
|
|
"pre_pad": self.pre_pad,
|
|
|
|
"scale": self.scale,
|
|
|
|
"tile_pad": self.tile_pad,
|
2023-02-02 04:31:01 +00:00
|
|
|
}
|