54 lines
1.3 KiB
Python
54 lines
1.3 KiB
Python
from typing import Optional
|
|
|
|
from ..params import (
|
|
HighresParams,
|
|
ImageParams,
|
|
Size,
|
|
SizeChart,
|
|
StageParams,
|
|
UpscaleParams,
|
|
)
|
|
from ..server import ServerContext
|
|
from ..worker import WorkerContext
|
|
from .base import BaseStage
|
|
from .result import StageResult
|
|
|
|
|
|
class EditMetadataStage(BaseStage):
|
|
max_tile = SizeChart.max
|
|
|
|
def run(
|
|
self,
|
|
_worker: WorkerContext,
|
|
_server: ServerContext,
|
|
_stage: StageParams,
|
|
_params: ImageParams,
|
|
source: StageResult,
|
|
*,
|
|
size: Optional[Size] = None,
|
|
upscale: Optional[UpscaleParams] = None,
|
|
highres: Optional[HighresParams] = None,
|
|
note: Optional[str] = None,
|
|
replace_params: Optional[ImageParams] = None,
|
|
**kwargs,
|
|
) -> StageResult:
|
|
# Modify the source image's metadata using the provided parameters
|
|
for metadata in source.metadata:
|
|
if note is not None:
|
|
metadata.note = note
|
|
|
|
if replace_params is not None:
|
|
metadata.params = replace_params
|
|
|
|
if size is not None:
|
|
metadata.size = size
|
|
|
|
if upscale is not None:
|
|
metadata.upscale = upscale
|
|
|
|
if highres is not None:
|
|
metadata.highres = highres
|
|
|
|
# Return the modified source image
|
|
return source
|