fix up highres metadata
This commit is contained in:
parent
ec48b27f93
commit
21666abf03
|
@ -102,7 +102,9 @@ class BlendImg2ImgStage(BaseStage):
|
||||||
|
|
||||||
outputs.extend(result.images)
|
outputs.extend(result.images)
|
||||||
|
|
||||||
metadata = [metadata.child(params) for metadata in sources.metadata]
|
metadata = [
|
||||||
|
metadata.child(params, metadata.size) for metadata in sources.metadata
|
||||||
|
]
|
||||||
return StageResult(images=outputs, metadata=metadata)
|
return StageResult(images=outputs, metadata=metadata)
|
||||||
|
|
||||||
def steps(
|
def steps(
|
||||||
|
|
|
@ -0,0 +1,39 @@
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from ..params import HighresParams, ImageParams, Size, StageParams, UpscaleParams
|
||||||
|
from ..server import ServerContext
|
||||||
|
from ..worker import WorkerContext
|
||||||
|
from .result import StageResult
|
||||||
|
|
||||||
|
|
||||||
|
class EditMetadataStage:
|
||||||
|
def run(
|
||||||
|
self,
|
||||||
|
_worker: WorkerContext,
|
||||||
|
_server: ServerContext,
|
||||||
|
_stage: StageParams,
|
||||||
|
_params: ImageParams,
|
||||||
|
source: StageResult,
|
||||||
|
*,
|
||||||
|
size: Size = None,
|
||||||
|
upscale: UpscaleParams = None,
|
||||||
|
highres: HighresParams = None,
|
||||||
|
note: str = None,
|
||||||
|
**kwargs,
|
||||||
|
) -> Any:
|
||||||
|
# Modify the source image's metadata using the provided parameters
|
||||||
|
for metadata in source.metadata:
|
||||||
|
if note is not None:
|
||||||
|
metadata.note = note
|
||||||
|
|
||||||
|
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
|
|
@ -2,6 +2,7 @@ from logging import getLogger
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
|
|
||||||
from ..chain.blend_img2img import BlendImg2ImgStage
|
from ..chain.blend_img2img import BlendImg2ImgStage
|
||||||
|
from ..chain.edit_metadata import EditMetadataStage
|
||||||
from ..chain.upscale import stage_upscale_correction
|
from ..chain.upscale import stage_upscale_correction
|
||||||
from ..chain.upscale_simple import UpscaleSimpleStage
|
from ..chain.upscale_simple import UpscaleSimpleStage
|
||||||
from ..params import HighresParams, ImageParams, StageParams, UpscaleParams
|
from ..params import HighresParams, ImageParams, StageParams, UpscaleParams
|
||||||
|
@ -66,4 +67,11 @@ def stage_highres(
|
||||||
strength=highres.strength,
|
strength=highres.strength,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# add highres parameters to the image metadata
|
||||||
|
chain.stage(
|
||||||
|
EditMetadataStage(),
|
||||||
|
stage.with_args(outscale=1),
|
||||||
|
highres=highres,
|
||||||
|
)
|
||||||
|
|
||||||
return chain
|
return chain
|
||||||
|
|
|
@ -28,6 +28,7 @@ class NetworkMetadata:
|
||||||
|
|
||||||
class ImageMetadata:
|
class ImageMetadata:
|
||||||
ancestors: List["ImageMetadata"]
|
ancestors: List["ImageMetadata"]
|
||||||
|
note: str
|
||||||
params: ImageParams
|
params: ImageParams
|
||||||
size: Size
|
size: Size
|
||||||
|
|
||||||
|
@ -70,6 +71,7 @@ class ImageMetadata:
|
||||||
self.loras = loras or []
|
self.loras = loras or []
|
||||||
self.models = models or []
|
self.models = models or []
|
||||||
self.ancestors = ancestors or []
|
self.ancestors = ancestors or []
|
||||||
|
self.note = ""
|
||||||
|
|
||||||
def child(
|
def child(
|
||||||
self,
|
self,
|
||||||
|
|
|
@ -9,6 +9,7 @@ from .blend_linear import BlendLinearStage
|
||||||
from .blend_mask import BlendMaskStage
|
from .blend_mask import BlendMaskStage
|
||||||
from .correct_codeformer import CorrectCodeformerStage
|
from .correct_codeformer import CorrectCodeformerStage
|
||||||
from .correct_gfpgan import CorrectGFPGANStage
|
from .correct_gfpgan import CorrectGFPGANStage
|
||||||
|
from .edit_metadata import EditMetadataStage
|
||||||
from .persist_disk import PersistDiskStage
|
from .persist_disk import PersistDiskStage
|
||||||
from .persist_s3 import PersistS3Stage
|
from .persist_s3 import PersistS3Stage
|
||||||
from .reduce_crop import ReduceCropStage
|
from .reduce_crop import ReduceCropStage
|
||||||
|
@ -38,6 +39,7 @@ CHAIN_STAGES = {
|
||||||
"blend-mask": BlendMaskStage,
|
"blend-mask": BlendMaskStage,
|
||||||
"correct-codeformer": CorrectCodeformerStage,
|
"correct-codeformer": CorrectCodeformerStage,
|
||||||
"correct-gfpgan": CorrectGFPGANStage,
|
"correct-gfpgan": CorrectGFPGANStage,
|
||||||
|
"edit-metadata": EditMetadataStage,
|
||||||
"persist-disk": PersistDiskStage,
|
"persist-disk": PersistDiskStage,
|
||||||
"persist-s3": PersistS3Stage,
|
"persist-s3": PersistS3Stage,
|
||||||
"reduce-crop": ReduceCropStage,
|
"reduce-crop": ReduceCropStage,
|
||||||
|
|
Loading…
Reference in New Issue