reset outscale for thumbnail metadata
This commit is contained in:
parent
95ac6d7c84
commit
5ffb44c8fa
|
@ -18,10 +18,12 @@ class EditSafetyStage(BaseStage):
|
|||
max_tile = SizeChart.max
|
||||
|
||||
def load(self, server: ServerContext) -> Any:
|
||||
# keep these within run to make this sort of like a plugin or peer dependency
|
||||
from horde_safety.deep_danbooru_model import get_deep_danbooru_model
|
||||
from horde_safety.interrogate import get_interrogator_no_blip
|
||||
from horde_safety.nsfw_checker_class import NSFWChecker
|
||||
|
||||
# check cache
|
||||
cache_key = ("horde-safety",)
|
||||
cache_checker = server.cache.get(ModelTypes.safety, cache_key)
|
||||
if cache_checker is not None:
|
||||
|
@ -53,7 +55,6 @@ class EditSafetyStage(BaseStage):
|
|||
) -> StageResult:
|
||||
logger.info("checking results using horde safety")
|
||||
|
||||
# keep these within run to make this sort of like a plugin or peer dependency
|
||||
try:
|
||||
# set up
|
||||
nsfw_checker = self.load(server)
|
||||
|
|
|
@ -91,6 +91,16 @@ def add_thumbnail_output(
|
|||
size=Size(server.thumbnail_size, server.thumbnail_size)
|
||||
)
|
||||
|
||||
if metadata.highres is not None:
|
||||
metadata.highres = metadata.highres.with_args(
|
||||
outscale=1,
|
||||
)
|
||||
|
||||
if metadata.upscale is not None:
|
||||
metadata.upscale = metadata.upscale.with_args(
|
||||
outscale=1,
|
||||
)
|
||||
|
||||
images.insert_image(0, thumbnail, metadata)
|
||||
|
||||
|
||||
|
|
|
@ -405,6 +405,9 @@ class StageParams:
|
|||
)
|
||||
|
||||
|
||||
UpscaleOrder = Literal["correction-first", "correction-last", "correction-both"]
|
||||
|
||||
|
||||
class UpscaleParams:
|
||||
def __init__(
|
||||
self,
|
||||
|
@ -419,9 +422,7 @@ class UpscaleParams:
|
|||
scale: int = 4,
|
||||
pre_pad: int = 0,
|
||||
tile_pad: int = 10,
|
||||
upscale_order: Literal[
|
||||
"correction-first", "correction-last", "correction-both"
|
||||
] = "correction-first",
|
||||
upscale_order: UpscaleOrder = "correction-first",
|
||||
) -> None:
|
||||
self.upscale_model = upscale_model
|
||||
self.correction_model = correction_model
|
||||
|
@ -478,23 +479,40 @@ class UpscaleParams:
|
|||
"upscale_order": self.upscale_order,
|
||||
}
|
||||
|
||||
def with_args(self, **kwargs):
|
||||
def with_args(
|
||||
self,
|
||||
upscale_model: Optional[str] = None,
|
||||
correction_model: Optional[str] = None,
|
||||
denoise: Optional[float] = None,
|
||||
upscale: Optional[bool] = None,
|
||||
faces: Optional[bool] = None,
|
||||
face_outscale: Optional[int] = None,
|
||||
face_strength: Optional[float] = None,
|
||||
outscale: Optional[int] = None,
|
||||
scale: Optional[int] = None,
|
||||
pre_pad: Optional[int] = None,
|
||||
tile_pad: Optional[int] = None,
|
||||
upscale_order: Optional[UpscaleOrder] = None,
|
||||
):
|
||||
return UpscaleParams(
|
||||
kwargs.get("upscale_model", self.upscale_model),
|
||||
kwargs.get("correction_model", self.correction_model),
|
||||
kwargs.get("denoise", self.denoise),
|
||||
kwargs.get("upscale", self.upscale),
|
||||
kwargs.get("faces", self.faces),
|
||||
kwargs.get("face_outscale", self.face_outscale),
|
||||
kwargs.get("face_strength", self.face_strength),
|
||||
kwargs.get("outscale", self.outscale),
|
||||
kwargs.get("scale", self.scale),
|
||||
kwargs.get("pre_pad", self.pre_pad),
|
||||
kwargs.get("tile_pad", self.tile_pad),
|
||||
kwargs.get("upscale_order", self.upscale_order),
|
||||
upscale_model=upscale_model or self.upscale_model,
|
||||
correction_model=correction_model or self.correction_model,
|
||||
denoise=denoise or self.denoise,
|
||||
upscale=upscale or self.upscale,
|
||||
faces=faces or self.faces,
|
||||
face_outscale=face_outscale or self.face_outscale,
|
||||
face_strength=face_strength or self.face_strength,
|
||||
outscale=outscale or self.outscale,
|
||||
scale=scale or self.scale,
|
||||
pre_pad=pre_pad or self.pre_pad,
|
||||
tile_pad=tile_pad or self.tile_pad,
|
||||
upscale_order=upscale_order or self.upscale_order,
|
||||
)
|
||||
|
||||
|
||||
UpscaleMethod = Literal["bilinear", "lanczos", "upscale"]
|
||||
|
||||
|
||||
class HighresParams:
|
||||
def __init__(
|
||||
self,
|
||||
|
@ -502,7 +520,7 @@ class HighresParams:
|
|||
scale: int,
|
||||
steps: int,
|
||||
strength: float,
|
||||
method: Literal["bilinear", "lanczos", "upscale"] = "lanczos",
|
||||
method: UpscaleMethod = "lanczos",
|
||||
iterations: int = 1,
|
||||
):
|
||||
self.enabled = enabled
|
||||
|
@ -531,3 +549,21 @@ class HighresParams:
|
|||
"steps": self.steps,
|
||||
"strength": self.strength,
|
||||
}
|
||||
|
||||
def with_args(
|
||||
self,
|
||||
enabled: Optional[bool] = None,
|
||||
scale: Optional[int] = None,
|
||||
steps: Optional[int] = None,
|
||||
strength: Optional[float] = None,
|
||||
method: Optional[UpscaleMethod] = None,
|
||||
iterations: Optional[int] = None,
|
||||
):
|
||||
return HighresParams(
|
||||
enabled=enabled or self.enabled,
|
||||
scale=scale or self.scale,
|
||||
steps=steps or self.steps,
|
||||
strength=strength or self.strength,
|
||||
method=method or self.method,
|
||||
iterations=iterations or self.iterations,
|
||||
)
|
||||
|
|
Loading…
Reference in New Issue