2023-01-28 23:09:19 +00:00
|
|
|
from logging import getLogger
|
2023-07-10 06:02:30 +00:00
|
|
|
from math import ceil
|
2023-06-30 04:06:36 +00:00
|
|
|
from typing import Any, List, Optional
|
2023-01-16 00:46:00 +00:00
|
|
|
|
2023-07-09 04:56:20 +00:00
|
|
|
from PIL import Image, ImageOps
|
2023-02-05 13:53:26 +00:00
|
|
|
|
2023-06-30 04:36:45 +00:00
|
|
|
from ..chain import (
|
2023-07-01 12:10:53 +00:00
|
|
|
BlendImg2ImgStage,
|
|
|
|
BlendMaskStage,
|
|
|
|
ChainPipeline,
|
|
|
|
SourceTxt2ImgStage,
|
|
|
|
UpscaleOutpaintStage,
|
2023-06-30 04:36:45 +00:00
|
|
|
)
|
2023-11-13 03:13:52 +00:00
|
|
|
from ..chain.highres import stage_highres
|
2024-01-06 20:17:26 +00:00
|
|
|
from ..chain.result import ImageMetadata, StageResult
|
2023-07-01 12:10:53 +00:00
|
|
|
from ..chain.upscale import split_upscale, stage_upscale_correction
|
2023-07-09 04:56:20 +00:00
|
|
|
from ..image import expand_image
|
2024-01-13 16:01:50 +00:00
|
|
|
from ..output import read_metadata, save_image, save_result
|
2023-04-01 17:06:31 +00:00
|
|
|
from ..params import (
|
|
|
|
Border,
|
|
|
|
HighresParams,
|
|
|
|
ImageParams,
|
|
|
|
Size,
|
|
|
|
StageParams,
|
|
|
|
UpscaleParams,
|
|
|
|
)
|
2023-02-26 05:49:39 +00:00
|
|
|
from ..server import ServerContext
|
2023-04-14 01:06:33 +00:00
|
|
|
from ..server.load import get_source_filters
|
2023-07-09 05:02:27 +00:00
|
|
|
from ..utils import is_debug, run_gc, show_system_toast
|
2023-02-26 20:15:30 +00:00
|
|
|
from ..worker import WorkerContext
|
2024-01-04 01:09:18 +00:00
|
|
|
from .utils import get_latents_from_seed
|
2023-01-28 05:28:14 +00:00
|
|
|
|
2023-01-28 23:09:19 +00:00
|
|
|
logger = getLogger(__name__)
|
|
|
|
|
2023-02-02 03:20:48 +00:00
|
|
|
|
2023-11-25 20:02:42 +00:00
|
|
|
def get_base_tile(params: ImageParams, size: Size) -> int:
|
|
|
|
if params.is_panorama():
|
|
|
|
tile = max(params.unet_tile, size.width, size.height)
|
|
|
|
logger.debug("adjusting tile size for panorama to %s", tile)
|
|
|
|
return tile
|
|
|
|
|
|
|
|
return params.unet_tile
|
|
|
|
|
|
|
|
|
|
|
|
def get_highres_tile(
|
|
|
|
server: ServerContext, params: ImageParams, highres: HighresParams, tile: int
|
|
|
|
) -> int:
|
|
|
|
if params.is_panorama() and server.has_feature("panorama-highres"):
|
|
|
|
return tile * highres.scale
|
|
|
|
|
|
|
|
return params.unet_tile
|
|
|
|
|
|
|
|
|
2024-01-09 03:37:44 +00:00
|
|
|
def add_safety_stage(
|
|
|
|
server: ServerContext,
|
|
|
|
pipeline: ChainPipeline,
|
|
|
|
) -> None:
|
|
|
|
if server.has_feature("horde-safety"):
|
|
|
|
from ..chain.edit_safety import EditSafetyStage
|
|
|
|
|
2024-01-09 04:14:32 +00:00
|
|
|
pipeline.stage(
|
|
|
|
EditSafetyStage(), StageParams(tile_size=EditSafetyStage.max_tile)
|
|
|
|
)
|
2024-01-09 03:37:44 +00:00
|
|
|
|
|
|
|
|
2024-01-13 16:01:50 +00:00
|
|
|
def add_thumbnail_output(
|
|
|
|
server: ServerContext,
|
|
|
|
images: StageResult,
|
|
|
|
params: ImageParams,
|
|
|
|
) -> None:
|
|
|
|
"""
|
|
|
|
Add a thumbnail image to the output, if requested.
|
|
|
|
TODO: This should really be a stage.
|
|
|
|
"""
|
|
|
|
result_size = images.size()
|
|
|
|
if (
|
|
|
|
params.thumbnail
|
|
|
|
and len(images) > 0
|
|
|
|
and (
|
|
|
|
result_size.width > server.thumbnail_size
|
|
|
|
or result_size.height > server.thumbnail_size
|
|
|
|
)
|
|
|
|
):
|
|
|
|
cover = images.as_images()[0]
|
|
|
|
thumbnail = cover.copy()
|
|
|
|
thumbnail.thumbnail((server.thumbnail_size, server.thumbnail_size))
|
|
|
|
|
2024-01-14 18:24:59 +00:00
|
|
|
metadata = images.metadata[0]
|
2024-01-14 18:48:24 +00:00
|
|
|
metadata = metadata.with_args(
|
2024-01-14 18:24:59 +00:00
|
|
|
size=Size(server.thumbnail_size, server.thumbnail_size)
|
|
|
|
)
|
|
|
|
|
2024-01-14 21:36:21 +00:00
|
|
|
if metadata.highres is not None:
|
|
|
|
metadata.highres = metadata.highres.with_args(
|
2024-01-15 02:45:28 +00:00
|
|
|
scale=1,
|
2024-01-14 21:36:21 +00:00
|
|
|
)
|
|
|
|
|
|
|
|
if metadata.upscale is not None:
|
|
|
|
metadata.upscale = metadata.upscale.with_args(
|
2024-01-15 02:45:28 +00:00
|
|
|
scale=1,
|
2024-01-14 21:36:21 +00:00
|
|
|
outscale=1,
|
|
|
|
)
|
|
|
|
|
2024-01-14 19:03:50 +00:00
|
|
|
images.insert_image(0, thumbnail, metadata)
|
2024-01-13 16:01:50 +00:00
|
|
|
|
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
def run_txt2img_pipeline(
|
2023-07-15 23:54:54 +00:00
|
|
|
worker: WorkerContext,
|
2023-04-14 13:54:21 +00:00
|
|
|
server: ServerContext,
|
|
|
|
params: ImageParams,
|
|
|
|
size: Size,
|
|
|
|
upscale: UpscaleParams,
|
|
|
|
highres: HighresParams,
|
2023-06-30 04:06:36 +00:00
|
|
|
) -> None:
|
2023-08-20 20:18:30 +00:00
|
|
|
# if using panorama, the pipeline will tile itself (views)
|
2023-11-25 20:02:42 +00:00
|
|
|
tile_size = get_base_tile(params, size)
|
2023-08-20 20:18:30 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# prepare the chain pipeline and first stage
|
|
|
|
chain = ChainPipeline()
|
2023-07-01 02:42:24 +00:00
|
|
|
chain.stage(
|
2023-07-01 12:10:53 +00:00
|
|
|
SourceTxt2ImgStage(),
|
2023-08-20 20:18:30 +00:00
|
|
|
StageParams(
|
|
|
|
tile_size=tile_size,
|
|
|
|
),
|
2023-07-01 02:42:24 +00:00
|
|
|
size=size,
|
2023-08-30 01:53:16 +00:00
|
|
|
prompt_index=0,
|
2023-11-05 01:41:58 +00:00
|
|
|
overlap=params.vae_overlap,
|
2023-06-30 04:42:52 +00:00
|
|
|
)
|
2023-06-30 04:06:36 +00:00
|
|
|
|
|
|
|
# apply upscaling and correction, before highres
|
2023-11-25 20:02:42 +00:00
|
|
|
highres_size = get_highres_tile(server, params, highres, tile_size)
|
2023-06-30 04:06:36 +00:00
|
|
|
first_upscale, after_upscale = split_upscale(upscale)
|
|
|
|
if first_upscale:
|
2023-07-01 02:42:24 +00:00
|
|
|
stage_upscale_correction(
|
2023-11-25 18:25:16 +00:00
|
|
|
StageParams(outscale=first_upscale.outscale, tile_size=highres_size),
|
2023-04-14 13:54:21 +00:00
|
|
|
params,
|
2023-06-30 04:06:36 +00:00
|
|
|
chain=chain,
|
2023-08-30 01:53:16 +00:00
|
|
|
upscale=first_upscale,
|
2023-04-14 13:54:21 +00:00
|
|
|
)
|
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# apply highres
|
2023-07-03 00:07:59 +00:00
|
|
|
stage_highres(
|
2023-11-25 18:25:16 +00:00
|
|
|
StageParams(outscale=highres.scale, tile_size=highres_size),
|
2023-07-03 00:07:59 +00:00
|
|
|
params,
|
|
|
|
highres,
|
|
|
|
upscale,
|
|
|
|
chain=chain,
|
2023-08-30 01:53:16 +00:00
|
|
|
prompt_index=1,
|
2023-07-03 00:07:59 +00:00
|
|
|
)
|
2023-04-27 12:22:00 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# apply upscaling and correction, after highres
|
2023-07-01 02:42:24 +00:00
|
|
|
stage_upscale_correction(
|
2023-11-25 18:25:16 +00:00
|
|
|
StageParams(outscale=after_upscale.outscale, tile_size=highres_size),
|
2023-05-02 04:20:40 +00:00
|
|
|
params,
|
2023-06-30 04:06:36 +00:00
|
|
|
chain=chain,
|
2023-08-30 01:53:16 +00:00
|
|
|
upscale=after_upscale,
|
2023-04-14 13:54:21 +00:00
|
|
|
)
|
|
|
|
|
2024-01-09 03:37:44 +00:00
|
|
|
add_safety_stage(server, chain)
|
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# run and save
|
2023-07-10 03:19:02 +00:00
|
|
|
latents = get_latents_from_seed(params.seed, size, batch=params.batch)
|
2024-01-04 03:39:19 +00:00
|
|
|
progress = worker.get_progress_callback(reset=True)
|
2024-01-04 01:09:18 +00:00
|
|
|
images = chain(
|
2023-11-20 00:39:39 +00:00
|
|
|
worker, server, params, StageResult.empty(), callback=progress, latents=latents
|
|
|
|
)
|
2023-04-14 13:54:21 +00:00
|
|
|
|
2024-01-13 16:01:50 +00:00
|
|
|
add_thumbnail_output(server, images, params)
|
2024-01-04 01:09:18 +00:00
|
|
|
save_result(server, images, worker.job)
|
2023-01-16 13:42:10 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# clean up
|
2023-07-15 23:54:54 +00:00
|
|
|
run_gc([worker.get_device()])
|
2023-06-30 04:06:36 +00:00
|
|
|
|
|
|
|
# notify the user
|
2024-01-04 01:09:18 +00:00
|
|
|
show_system_toast(f"finished txt2img job: {worker.job}")
|
|
|
|
logger.info("finished txt2img job: %s", worker.job)
|
2023-01-16 00:54:20 +00:00
|
|
|
|
|
|
|
|
2023-01-16 13:31:42 +00:00
|
|
|
def run_img2img_pipeline(
|
2023-07-15 23:54:54 +00:00
|
|
|
worker: WorkerContext,
|
2023-02-04 16:06:22 +00:00
|
|
|
server: ServerContext,
|
2023-01-27 23:08:36 +00:00
|
|
|
params: ImageParams,
|
2023-01-16 19:12:08 +00:00
|
|
|
upscale: UpscaleParams,
|
2023-04-14 13:54:21 +00:00
|
|
|
highres: HighresParams,
|
2023-02-18 22:35:57 +00:00
|
|
|
source: Image.Image,
|
2023-01-16 19:12:08 +00:00
|
|
|
strength: float,
|
2023-04-14 01:06:33 +00:00
|
|
|
source_filter: Optional[str] = None,
|
2023-01-28 18:42:02 +00:00
|
|
|
) -> None:
|
2023-06-30 04:06:36 +00:00
|
|
|
# run filter on the source image
|
2023-04-14 01:06:33 +00:00
|
|
|
if source_filter is not None:
|
2023-04-14 02:10:00 +00:00
|
|
|
f = get_source_filters().get(source_filter, None)
|
|
|
|
if f is not None:
|
2023-04-22 16:26:21 +00:00
|
|
|
logger.debug("running source filter: %s", f.__name__)
|
2023-04-14 02:10:00 +00:00
|
|
|
source = f(server, source)
|
2023-04-14 01:06:33 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# prepare the chain pipeline and first stage
|
2023-11-25 20:02:42 +00:00
|
|
|
tile_size = get_base_tile(params, Size(*source.size))
|
2023-06-30 04:06:36 +00:00
|
|
|
chain = ChainPipeline()
|
2023-07-01 02:42:24 +00:00
|
|
|
chain.stage(
|
2023-07-01 12:10:53 +00:00
|
|
|
BlendImg2ImgStage(),
|
2023-11-25 20:02:42 +00:00
|
|
|
StageParams(
|
|
|
|
tile_size=tile_size,
|
|
|
|
),
|
2023-08-30 01:53:16 +00:00
|
|
|
prompt_index=0,
|
2023-07-01 02:42:24 +00:00
|
|
|
strength=strength,
|
2023-11-05 01:41:58 +00:00
|
|
|
overlap=params.vae_overlap,
|
2023-02-05 13:53:26 +00:00
|
|
|
)
|
2023-04-13 04:30:59 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# apply upscaling and correction, before highres
|
|
|
|
first_upscale, after_upscale = split_upscale(upscale)
|
|
|
|
if first_upscale:
|
2023-07-01 02:42:24 +00:00
|
|
|
stage_upscale_correction(
|
2023-11-25 20:02:42 +00:00
|
|
|
StageParams(
|
|
|
|
outscale=first_upscale.outscale,
|
|
|
|
tile_size=tile_size,
|
|
|
|
),
|
2023-06-30 04:06:36 +00:00
|
|
|
params,
|
|
|
|
upscale=first_upscale,
|
|
|
|
chain=chain,
|
2023-02-05 23:36:00 +00:00
|
|
|
)
|
2023-01-16 00:54:20 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# loopback through multiple img2img iterations
|
2023-07-02 17:16:13 +00:00
|
|
|
for _i in range(params.loopback):
|
|
|
|
chain.stage(
|
|
|
|
BlendImg2ImgStage(),
|
2023-11-25 20:02:42 +00:00
|
|
|
StageParams(
|
|
|
|
tile_size=tile_size,
|
|
|
|
),
|
2023-07-02 17:16:13 +00:00
|
|
|
strength=strength,
|
|
|
|
)
|
2023-04-14 03:51:59 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# highres, if selected
|
2023-11-25 20:02:42 +00:00
|
|
|
highres_size = get_highres_tile(server, params, highres, tile_size)
|
2023-07-03 00:07:59 +00:00
|
|
|
stage_highres(
|
2023-11-25 20:02:42 +00:00
|
|
|
StageParams(tile_size=highres_size, outscale=highres.scale),
|
2023-07-03 00:07:59 +00:00
|
|
|
params,
|
|
|
|
highres,
|
|
|
|
upscale,
|
|
|
|
chain=chain,
|
2023-08-30 01:53:16 +00:00
|
|
|
prompt_index=1,
|
2023-07-03 00:07:59 +00:00
|
|
|
)
|
2023-04-22 15:54:39 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# apply upscaling and correction, after highres
|
2023-07-01 02:42:24 +00:00
|
|
|
stage_upscale_correction(
|
2023-11-25 20:02:42 +00:00
|
|
|
StageParams(tile_size=tile_size, outscale=after_upscale.scale),
|
2023-06-30 04:06:36 +00:00
|
|
|
params,
|
|
|
|
upscale=after_upscale,
|
|
|
|
chain=chain,
|
|
|
|
)
|
2023-04-14 13:54:21 +00:00
|
|
|
|
2024-01-09 03:37:44 +00:00
|
|
|
add_safety_stage(server, chain)
|
|
|
|
|
2024-01-13 16:01:50 +00:00
|
|
|
# prep inputs
|
|
|
|
input_metadata = read_metadata(source) or ImageMetadata.unknown_image()
|
|
|
|
input_result = StageResult(images=[source], metadata=[input_metadata])
|
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# run and append the filtered source
|
2024-01-04 03:39:19 +00:00
|
|
|
progress = worker.get_progress_callback(reset=True)
|
2024-01-04 01:09:18 +00:00
|
|
|
images = chain(
|
2024-01-06 20:17:26 +00:00
|
|
|
worker,
|
|
|
|
server,
|
|
|
|
params,
|
2024-01-13 16:01:50 +00:00
|
|
|
input_result, # terrible naming, I know
|
2024-01-06 20:17:26 +00:00
|
|
|
callback=progress,
|
2023-11-20 00:39:39 +00:00
|
|
|
)
|
2023-06-30 04:06:36 +00:00
|
|
|
|
|
|
|
if source_filter is not None and source_filter != "none":
|
2024-01-09 15:21:51 +00:00
|
|
|
images.push_image(source, ImageMetadata.unknown_image())
|
2023-06-30 04:06:36 +00:00
|
|
|
|
2024-01-13 16:01:50 +00:00
|
|
|
add_thumbnail_output(server, images, params)
|
2024-01-04 01:09:18 +00:00
|
|
|
save_result(server, images, worker.job)
|
2023-01-16 13:42:10 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# clean up
|
2023-07-15 23:54:54 +00:00
|
|
|
run_gc([worker.get_device()])
|
2023-06-30 04:06:36 +00:00
|
|
|
|
|
|
|
# notify the user
|
2024-01-04 01:09:18 +00:00
|
|
|
show_system_toast(f"finished img2img job: {worker.job}")
|
|
|
|
logger.info("finished img2img job: %s", worker.job)
|
2023-01-16 00:54:20 +00:00
|
|
|
|
|
|
|
|
|
|
|
def run_inpaint_pipeline(
|
2023-07-15 23:54:54 +00:00
|
|
|
worker: WorkerContext,
|
2023-02-04 16:06:22 +00:00
|
|
|
server: ServerContext,
|
2023-01-27 23:08:36 +00:00
|
|
|
params: ImageParams,
|
2023-02-02 04:20:40 +00:00
|
|
|
size: Size,
|
2023-01-16 19:12:08 +00:00
|
|
|
upscale: UpscaleParams,
|
2023-04-14 13:54:21 +00:00
|
|
|
highres: HighresParams,
|
2023-02-18 22:35:57 +00:00
|
|
|
source: Image.Image,
|
|
|
|
mask: Image.Image,
|
2023-02-02 04:20:40 +00:00
|
|
|
border: Border,
|
2023-01-16 00:54:20 +00:00
|
|
|
noise_source: Any,
|
2023-01-17 23:50:36 +00:00
|
|
|
mask_filter: Any,
|
2023-01-18 14:41:02 +00:00
|
|
|
fill_color: str,
|
2023-02-12 00:00:18 +00:00
|
|
|
tile_order: str,
|
2023-07-11 03:16:17 +00:00
|
|
|
full_res_inpaint: bool,
|
|
|
|
full_res_inpaint_padding: float,
|
2023-01-28 18:42:02 +00:00
|
|
|
) -> None:
|
2023-06-30 04:06:36 +00:00
|
|
|
logger.debug("building inpaint pipeline")
|
2023-11-25 20:02:42 +00:00
|
|
|
tile_size = get_base_tile(params, size)
|
2023-07-09 05:02:27 +00:00
|
|
|
|
2023-07-09 04:56:20 +00:00
|
|
|
if mask is None:
|
|
|
|
# if no mask was provided, keep the full source image
|
|
|
|
mask = Image.new("L", source.size, 0)
|
|
|
|
|
|
|
|
# masks start as 512x512, resize to cover the source, then trim the extra
|
|
|
|
mask_max = max(source.width, source.height)
|
|
|
|
mask = ImageOps.contain(mask, (mask_max, mask_max))
|
|
|
|
mask = mask.crop((0, 0, source.width, source.height))
|
|
|
|
|
2024-01-04 01:09:18 +00:00
|
|
|
source, mask, noise, _full_size = expand_image(
|
2023-07-09 04:56:20 +00:00
|
|
|
source,
|
|
|
|
mask,
|
|
|
|
border,
|
|
|
|
fill=fill_color,
|
|
|
|
noise_source=noise_source,
|
|
|
|
mask_filter=mask_filter,
|
|
|
|
)
|
|
|
|
|
2023-07-11 02:21:10 +00:00
|
|
|
if is_debug():
|
|
|
|
save_image(server, "full-source.png", source)
|
|
|
|
save_image(server, "full-mask.png", mask)
|
|
|
|
save_image(server, "full-noise.png", noise)
|
|
|
|
|
2023-07-10 06:02:30 +00:00
|
|
|
# check if we can do full-res inpainting if no outpainting is done
|
|
|
|
logger.debug("border zero: %s", border.isZero())
|
2023-07-13 00:42:04 +00:00
|
|
|
full_res_inpaint = full_res_inpaint and border.isZero()
|
|
|
|
if full_res_inpaint:
|
2023-11-25 21:25:47 +00:00
|
|
|
bbox = mask.getbbox()
|
|
|
|
if bbox is None:
|
|
|
|
bbox = (0, 0, source.width, source.height)
|
|
|
|
|
|
|
|
logger.debug("mask bounding box: %s", bbox)
|
|
|
|
mask_left, mask_top, mask_right, mask_bottom = bbox
|
2023-07-10 06:02:30 +00:00
|
|
|
mask_width = mask_right - mask_left
|
|
|
|
mask_height = mask_bottom - mask_top
|
|
|
|
# ensure we have some padding around the mask when we do the inpaint (and that the region size is even)
|
2023-07-11 03:24:36 +00:00
|
|
|
adj_mask_size = (
|
|
|
|
ceil(max(mask_width, mask_height) * full_res_inpaint_padding / 2) * 2
|
|
|
|
)
|
2023-07-10 06:15:46 +00:00
|
|
|
mask_center_x = int(round((mask_right + mask_left) / 2))
|
|
|
|
mask_center_y = int(round((mask_bottom + mask_top) / 2))
|
|
|
|
adj_mask_border = (
|
|
|
|
int(mask_center_x - adj_mask_size / 2),
|
|
|
|
int(mask_center_y - adj_mask_size / 2),
|
|
|
|
int(mask_center_x + adj_mask_size / 2),
|
|
|
|
int(mask_center_y + adj_mask_size / 2),
|
|
|
|
)
|
2023-07-11 03:24:36 +00:00
|
|
|
|
|
|
|
# we would like to subtract the excess width (subtract a positive) and add the deficient width (subtract a negative)
|
|
|
|
x_adj = -max(adj_mask_border[2] - source.width, 0) - min(adj_mask_border[0], 0)
|
|
|
|
# we would like to subtract the excess height (subtract a negative) and add the deficient height (subtract a negative)
|
|
|
|
y_adj = -max(adj_mask_border[3] - source.height, 0) - min(adj_mask_border[1], 0)
|
|
|
|
|
|
|
|
adj_mask_border = (
|
|
|
|
adj_mask_border[0] + x_adj,
|
|
|
|
adj_mask_border[1] + y_adj,
|
|
|
|
adj_mask_border[2] + x_adj,
|
|
|
|
adj_mask_border[3] + y_adj,
|
|
|
|
)
|
2023-07-10 14:24:39 +00:00
|
|
|
|
2023-07-10 06:15:46 +00:00
|
|
|
border_integrity = all(
|
2023-07-10 06:32:02 +00:00
|
|
|
(
|
|
|
|
adj_mask_border[0] >= 0,
|
|
|
|
adj_mask_border[1] >= 0,
|
|
|
|
adj_mask_border[2] <= source.width,
|
|
|
|
adj_mask_border[3] <= source.height,
|
|
|
|
)
|
2023-07-10 06:15:46 +00:00
|
|
|
)
|
|
|
|
logger.debug(
|
|
|
|
"adjusted mask size %s, mask bounding box: %s",
|
|
|
|
adj_mask_size,
|
|
|
|
adj_mask_border,
|
|
|
|
)
|
2023-07-10 06:33:13 +00:00
|
|
|
if border_integrity and adj_mask_size <= tile_size:
|
2023-07-11 03:16:17 +00:00
|
|
|
logger.debug("performing full-res inpainting")
|
2023-07-10 06:02:30 +00:00
|
|
|
original_source = source
|
|
|
|
source = source.crop(adj_mask_border)
|
|
|
|
source = ImageOps.contain(source, (tile_size, tile_size))
|
|
|
|
mask = mask.crop(adj_mask_border)
|
|
|
|
mask = ImageOps.contain(mask, (tile_size, tile_size))
|
|
|
|
if is_debug():
|
|
|
|
save_image(server, "adjusted-mask.png", mask)
|
|
|
|
save_image(server, "adjusted-source.png", source)
|
2023-07-11 03:16:17 +00:00
|
|
|
else:
|
|
|
|
logger.debug("cannot perform full-res inpaint due to size issue")
|
|
|
|
full_res_inpaint = False
|
2023-07-10 06:02:30 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# set up the chain pipeline and base stage
|
|
|
|
chain = ChainPipeline()
|
2023-07-01 02:42:24 +00:00
|
|
|
chain.stage(
|
2023-07-01 12:10:53 +00:00
|
|
|
UpscaleOutpaintStage(),
|
2023-11-25 20:02:42 +00:00
|
|
|
StageParams(tile_order=tile_order, tile_size=tile_size),
|
2023-07-01 02:42:24 +00:00
|
|
|
border=border,
|
2023-07-09 04:56:20 +00:00
|
|
|
mask=mask,
|
2023-07-01 02:42:24 +00:00
|
|
|
fill_color=fill_color,
|
|
|
|
mask_filter=mask_filter,
|
|
|
|
noise_source=noise_source,
|
2023-11-05 01:41:58 +00:00
|
|
|
overlap=params.vae_overlap,
|
2023-08-30 01:53:16 +00:00
|
|
|
prompt_index=0,
|
2023-01-22 04:28:13 +00:00
|
|
|
)
|
2023-01-16 00:54:20 +00:00
|
|
|
|
2023-07-03 00:07:59 +00:00
|
|
|
# apply upscaling and correction, before highres
|
|
|
|
first_upscale, after_upscale = split_upscale(upscale)
|
|
|
|
if first_upscale:
|
|
|
|
stage_upscale_correction(
|
2023-11-25 20:02:42 +00:00
|
|
|
StageParams(outscale=first_upscale.outscale, tile_size=tile_size),
|
2023-07-02 17:16:13 +00:00
|
|
|
params,
|
2023-07-03 00:07:59 +00:00
|
|
|
upscale=first_upscale,
|
2023-07-02 17:16:13 +00:00
|
|
|
chain=chain,
|
|
|
|
)
|
2023-04-15 01:29:44 +00:00
|
|
|
|
2023-07-03 00:07:59 +00:00
|
|
|
# apply highres
|
2023-11-25 20:02:42 +00:00
|
|
|
highres_size = get_highres_tile(server, params, highres, tile_size)
|
2023-07-03 00:07:59 +00:00
|
|
|
stage_highres(
|
2023-11-25 20:02:42 +00:00
|
|
|
StageParams(outscale=highres.scale, tile_size=highres_size),
|
2023-07-03 00:07:59 +00:00
|
|
|
params,
|
|
|
|
highres,
|
|
|
|
upscale,
|
|
|
|
chain=chain,
|
2023-08-30 01:53:16 +00:00
|
|
|
prompt_index=1,
|
2023-07-03 00:07:59 +00:00
|
|
|
)
|
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# apply upscaling and correction
|
2023-07-01 02:42:24 +00:00
|
|
|
stage_upscale_correction(
|
2023-11-25 20:02:42 +00:00
|
|
|
StageParams(outscale=after_upscale.outscale),
|
2023-04-12 00:29:25 +00:00
|
|
|
params,
|
2023-07-03 00:07:59 +00:00
|
|
|
upscale=after_upscale,
|
2023-06-30 04:06:36 +00:00
|
|
|
chain=chain,
|
2023-02-12 18:33:36 +00:00
|
|
|
)
|
2023-01-16 19:12:08 +00:00
|
|
|
|
2024-01-09 03:37:44 +00:00
|
|
|
add_safety_stage(server, chain)
|
|
|
|
|
2024-01-13 16:01:50 +00:00
|
|
|
# prep inputs
|
|
|
|
input_metadata = read_metadata(source) or ImageMetadata.unknown_image()
|
|
|
|
input_result = StageResult(images=[source], metadata=[input_metadata])
|
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# run and save
|
2023-07-10 03:19:02 +00:00
|
|
|
latents = get_latents_from_seed(params.seed, size, batch=params.batch)
|
2024-01-04 03:39:19 +00:00
|
|
|
progress = worker.get_progress_callback(reset=True)
|
2024-01-04 01:09:18 +00:00
|
|
|
images = chain(
|
2023-11-24 16:36:53 +00:00
|
|
|
worker,
|
|
|
|
server,
|
|
|
|
params,
|
2024-01-13 16:01:50 +00:00
|
|
|
input_result,
|
2023-11-24 16:36:53 +00:00
|
|
|
callback=progress,
|
|
|
|
latents=latents,
|
2023-11-20 00:39:39 +00:00
|
|
|
)
|
2023-06-30 04:06:36 +00:00
|
|
|
|
2024-01-13 16:01:50 +00:00
|
|
|
add_thumbnail_output(server, images, params)
|
|
|
|
|
2024-01-06 20:27:55 +00:00
|
|
|
for i, (image, metadata) in enumerate(zip(images.as_images(), images.metadata)):
|
2023-07-10 06:02:30 +00:00
|
|
|
if full_res_inpaint:
|
|
|
|
if is_debug():
|
|
|
|
save_image(server, "adjusted-output.png", image)
|
2024-01-04 01:09:18 +00:00
|
|
|
|
2023-07-10 06:02:30 +00:00
|
|
|
mini_image = ImageOps.contain(image, (adj_mask_size, adj_mask_size))
|
|
|
|
image = original_source
|
|
|
|
image.paste(mini_image, box=adj_mask_border)
|
2024-01-04 01:09:18 +00:00
|
|
|
|
|
|
|
save_image(
|
2023-07-04 18:56:02 +00:00
|
|
|
server,
|
2024-01-04 01:09:18 +00:00
|
|
|
f"{worker.job}_{i}.{server.image_format}",
|
2023-07-04 18:56:02 +00:00
|
|
|
image,
|
2024-01-04 01:09:18 +00:00
|
|
|
metadata,
|
2023-07-04 18:56:02 +00:00
|
|
|
)
|
2023-01-16 00:54:20 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# clean up
|
2023-07-15 23:54:54 +00:00
|
|
|
run_gc([worker.get_device()])
|
2023-06-26 12:03:06 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# notify the user
|
2024-01-04 01:09:18 +00:00
|
|
|
show_system_toast(f"finished inpaint job: {worker.job}")
|
|
|
|
logger.info("finished inpaint job: %s", worker.job)
|
2023-01-17 05:45:54 +00:00
|
|
|
|
2023-01-17 23:50:36 +00:00
|
|
|
|
2023-01-17 05:45:54 +00:00
|
|
|
def run_upscale_pipeline(
|
2023-07-15 23:54:54 +00:00
|
|
|
worker: WorkerContext,
|
2023-02-04 16:06:22 +00:00
|
|
|
server: ServerContext,
|
2023-01-27 23:08:36 +00:00
|
|
|
params: ImageParams,
|
2023-02-02 04:20:40 +00:00
|
|
|
size: Size,
|
2023-01-17 05:45:54 +00:00
|
|
|
upscale: UpscaleParams,
|
2023-04-14 13:54:21 +00:00
|
|
|
highres: HighresParams,
|
2023-02-18 22:35:57 +00:00
|
|
|
source: Image.Image,
|
2023-01-28 18:42:02 +00:00
|
|
|
) -> None:
|
2023-06-30 04:06:36 +00:00
|
|
|
# set up the chain pipeline, no base stage for upscaling
|
|
|
|
chain = ChainPipeline()
|
2023-11-25 20:02:42 +00:00
|
|
|
tile_size = get_base_tile(params, size)
|
2023-02-04 16:06:22 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# apply upscaling and correction, before highres
|
|
|
|
first_upscale, after_upscale = split_upscale(upscale)
|
|
|
|
if first_upscale:
|
2023-07-01 02:42:24 +00:00
|
|
|
stage_upscale_correction(
|
2023-11-25 20:02:42 +00:00
|
|
|
StageParams(outscale=first_upscale.outscale, tile_size=tile_size),
|
2023-06-30 04:06:36 +00:00
|
|
|
params,
|
|
|
|
upscale=first_upscale,
|
|
|
|
chain=chain,
|
|
|
|
)
|
2023-04-15 01:29:44 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# apply highres
|
2023-11-25 20:02:42 +00:00
|
|
|
highres_size = get_highres_tile(server, params, highres, tile_size)
|
2023-07-03 00:07:59 +00:00
|
|
|
stage_highres(
|
2023-11-25 20:02:42 +00:00
|
|
|
StageParams(outscale=highres.scale, tile_size=highres_size),
|
2023-07-03 00:07:59 +00:00
|
|
|
params,
|
|
|
|
highres,
|
|
|
|
upscale,
|
|
|
|
chain=chain,
|
2023-08-30 01:53:16 +00:00
|
|
|
prompt_index=0,
|
2023-07-03 00:07:59 +00:00
|
|
|
)
|
2023-06-30 04:06:36 +00:00
|
|
|
|
|
|
|
# apply upscaling and correction, after highres
|
2023-07-01 02:42:24 +00:00
|
|
|
stage_upscale_correction(
|
2023-11-25 20:02:42 +00:00
|
|
|
StageParams(outscale=after_upscale.outscale, tile_size=tile_size),
|
2023-06-30 04:06:36 +00:00
|
|
|
params,
|
|
|
|
upscale=after_upscale,
|
|
|
|
chain=chain,
|
2023-02-05 13:53:26 +00:00
|
|
|
)
|
2023-01-17 05:45:54 +00:00
|
|
|
|
2024-01-09 03:37:44 +00:00
|
|
|
add_safety_stage(server, chain)
|
|
|
|
|
2024-01-13 16:01:50 +00:00
|
|
|
# prep inputs
|
|
|
|
input_metadata = read_metadata(source) or ImageMetadata.unknown_image()
|
|
|
|
input_result = StageResult(images=[source], metadata=[input_metadata])
|
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# run and save
|
2024-01-04 03:39:19 +00:00
|
|
|
progress = worker.get_progress_callback(reset=True)
|
2024-01-04 01:09:18 +00:00
|
|
|
images = chain(
|
2024-01-06 20:17:26 +00:00
|
|
|
worker,
|
|
|
|
server,
|
|
|
|
params,
|
2024-01-13 16:01:50 +00:00
|
|
|
input_result,
|
2024-01-06 20:17:26 +00:00
|
|
|
callback=progress,
|
2023-11-20 00:39:39 +00:00
|
|
|
)
|
2023-07-01 03:00:11 +00:00
|
|
|
|
2024-01-13 16:01:50 +00:00
|
|
|
add_thumbnail_output(server, images, params)
|
2024-01-04 01:09:18 +00:00
|
|
|
save_result(server, images, worker.job)
|
2023-04-15 01:29:44 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# clean up
|
2023-07-15 23:54:54 +00:00
|
|
|
run_gc([worker.get_device()])
|
2023-06-26 12:03:06 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# notify the user
|
2024-01-04 01:09:18 +00:00
|
|
|
show_system_toast(f"finished upscale job: {worker.job}")
|
|
|
|
logger.info("finished upscale job: %s", worker.job)
|
2023-02-13 23:34:42 +00:00
|
|
|
|
|
|
|
|
|
|
|
def run_blend_pipeline(
|
2023-07-15 23:54:54 +00:00
|
|
|
worker: WorkerContext,
|
2023-02-13 23:34:42 +00:00
|
|
|
server: ServerContext,
|
|
|
|
params: ImageParams,
|
|
|
|
size: Size,
|
|
|
|
upscale: UpscaleParams,
|
2023-04-22 17:40:51 +00:00
|
|
|
# highres: HighresParams,
|
2023-11-19 03:35:00 +00:00
|
|
|
sources: List[Image.Image],
|
2023-02-13 23:34:42 +00:00
|
|
|
mask: Image.Image,
|
|
|
|
) -> None:
|
2023-06-30 04:06:36 +00:00
|
|
|
# set up the chain pipeline and base stage
|
|
|
|
chain = ChainPipeline()
|
2023-11-25 20:02:42 +00:00
|
|
|
tile_size = get_base_tile(params, size)
|
|
|
|
|
2023-12-21 05:41:54 +00:00
|
|
|
# resize mask to match source size
|
2024-01-06 22:27:25 +00:00
|
|
|
stage_source = sources.pop()
|
2023-12-21 05:41:54 +00:00
|
|
|
stage_mask = mask.resize(stage_source.size, Image.Resampling.BILINEAR)
|
|
|
|
|
2023-11-25 20:02:42 +00:00
|
|
|
chain.stage(
|
|
|
|
BlendMaskStage(),
|
|
|
|
StageParams(tile_size=tile_size),
|
2023-12-21 05:41:54 +00:00
|
|
|
stage_source=stage_source,
|
|
|
|
stage_mask=stage_mask,
|
2023-11-25 20:02:42 +00:00
|
|
|
)
|
2023-02-13 23:34:42 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# apply upscaling and correction
|
2023-07-01 02:42:24 +00:00
|
|
|
stage_upscale_correction(
|
2023-11-25 20:02:42 +00:00
|
|
|
StageParams(outscale=upscale.outscale),
|
2023-02-13 23:34:42 +00:00
|
|
|
params,
|
2023-06-30 04:06:36 +00:00
|
|
|
upscale=upscale,
|
|
|
|
chain=chain,
|
2023-02-13 23:34:42 +00:00
|
|
|
)
|
|
|
|
|
2024-01-09 03:37:44 +00:00
|
|
|
add_safety_stage(server, chain)
|
|
|
|
|
2024-01-13 16:01:50 +00:00
|
|
|
# prep inputs
|
|
|
|
input_metadata = [
|
|
|
|
read_metadata(source) or ImageMetadata.unknown_image() for source in sources
|
|
|
|
]
|
|
|
|
input_result = StageResult(images=sources, metadata=input_metadata)
|
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# run and save
|
2024-01-04 03:39:19 +00:00
|
|
|
progress = worker.get_progress_callback(reset=True)
|
2024-01-04 01:09:18 +00:00
|
|
|
images = chain(
|
2024-01-06 20:17:26 +00:00
|
|
|
worker,
|
|
|
|
server,
|
|
|
|
params,
|
2024-01-13 16:01:50 +00:00
|
|
|
input_result,
|
2024-01-06 20:17:26 +00:00
|
|
|
callback=progress,
|
2023-11-20 00:39:39 +00:00
|
|
|
)
|
2023-07-04 18:56:02 +00:00
|
|
|
|
2024-01-13 16:01:50 +00:00
|
|
|
add_thumbnail_output(server, images, params)
|
2024-01-04 01:09:18 +00:00
|
|
|
save_result(server, images, worker.job)
|
2023-02-13 23:34:42 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# clean up
|
2023-07-15 23:54:54 +00:00
|
|
|
run_gc([worker.get_device()])
|
2023-06-26 12:03:06 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# notify the user
|
2024-01-04 01:09:18 +00:00
|
|
|
show_system_toast(f"finished blend job: {worker.job}")
|
|
|
|
logger.info("finished blend job: %s", worker.job)
|