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-11-13 03:13:52 +00:00
|
|
|
BlendDenoiseStage,
|
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
|
2023-11-19 00:13:13 +00:00
|
|
|
from ..chain.result import 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
|
2023-06-26 12:14:32 +00:00
|
|
|
from ..output import save_image
|
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
|
2023-07-10 03:19:02 +00:00
|
|
|
from .utils import get_latents_from_seed, parse_prompt
|
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
|
|
|
|
|
|
|
|
|
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,
|
2023-06-30 04:06:36 +00:00
|
|
|
outputs: List[str],
|
2023-04-14 13:54:21 +00:00
|
|
|
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-11-13 03:19:01 +00:00
|
|
|
if params.is_panorama():
|
|
|
|
chain.stage(
|
|
|
|
BlendDenoiseStage(),
|
2023-11-25 18:25:16 +00:00
|
|
|
StageParams(tile_size=highres_size),
|
2023-11-13 03:19:01 +00:00
|
|
|
)
|
2023-11-11 23:03:01 +00:00
|
|
|
|
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
|
|
|
)
|
|
|
|
|
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)
|
2023-07-15 23:54:54 +00:00
|
|
|
progress = worker.get_progress_callback()
|
2023-11-20 00:39:39 +00:00
|
|
|
images = chain.run(
|
|
|
|
worker, server, params, StageResult.empty(), callback=progress, latents=latents
|
|
|
|
)
|
2023-04-14 13:54:21 +00:00
|
|
|
|
2023-07-07 01:39:08 +00:00
|
|
|
_pairs, loras, inversions, _rest = parse_prompt(params)
|
2023-07-04 18:56:02 +00:00
|
|
|
|
|
|
|
for image, output in zip(images, outputs):
|
2023-09-28 23:45:04 +00:00
|
|
|
logger.trace("saving output image %s: %s", output, image.size)
|
2023-07-04 18:56:02 +00:00
|
|
|
dest = save_image(
|
|
|
|
server,
|
|
|
|
output,
|
|
|
|
image,
|
|
|
|
params,
|
|
|
|
size,
|
|
|
|
upscale=upscale,
|
|
|
|
highres=highres,
|
|
|
|
inversions=inversions,
|
|
|
|
loras=loras,
|
|
|
|
)
|
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
|
2023-04-24 22:40:50 +00:00
|
|
|
show_system_toast(f"finished txt2img job: {dest}")
|
2023-02-05 13:53:26 +00:00
|
|
|
logger.info("finished txt2img job: %s", dest)
|
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-02-20 14:35:18 +00:00
|
|
|
outputs: List[str],
|
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
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# run and append the filtered source
|
2023-07-15 23:54:54 +00:00
|
|
|
progress = worker.get_progress_callback()
|
2023-11-20 00:39:39 +00:00
|
|
|
images = chain.run(
|
|
|
|
worker, server, params, StageResult(images=[source]), callback=progress
|
|
|
|
)
|
2023-06-30 04:06:36 +00:00
|
|
|
|
|
|
|
if source_filter is not None and source_filter != "none":
|
|
|
|
images.append(source)
|
|
|
|
|
|
|
|
# save with metadata
|
2023-07-07 01:39:08 +00:00
|
|
|
_pairs, loras, inversions, _rest = parse_prompt(params)
|
2023-06-30 04:06:36 +00:00
|
|
|
size = Size(*source.size)
|
2023-01-16 00:54:20 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
for image, output in zip(images, outputs):
|
2023-06-26 12:14:32 +00:00
|
|
|
dest = save_image(
|
2023-06-26 12:48:39 +00:00
|
|
|
server,
|
|
|
|
output,
|
|
|
|
image,
|
|
|
|
params,
|
|
|
|
size,
|
|
|
|
upscale=upscale,
|
|
|
|
highres=highres,
|
|
|
|
inversions=inversions,
|
|
|
|
loras=loras,
|
2023-06-26 12:14:32 +00:00
|
|
|
)
|
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
|
2023-04-24 22:40:50 +00:00
|
|
|
show_system_toast(f"finished img2img job: {dest}")
|
2023-02-05 13:53:26 +00:00
|
|
|
logger.info("finished img2img job: %s", dest)
|
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-02-20 14:35:18 +00:00
|
|
|
outputs: List[str],
|
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))
|
|
|
|
|
|
|
|
source, mask, noise, full_size = expand_image(
|
|
|
|
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
|
|
|
|
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)
|
2023-07-15 23:54:54 +00:00
|
|
|
progress = worker.get_progress_callback()
|
2023-11-20 00:39:39 +00:00
|
|
|
images = chain.run(
|
2023-11-24 16:36:53 +00:00
|
|
|
worker,
|
|
|
|
server,
|
|
|
|
params,
|
|
|
|
StageResult(images=[source]),
|
|
|
|
callback=progress,
|
|
|
|
latents=latents,
|
2023-11-20 00:39:39 +00:00
|
|
|
)
|
2023-06-30 04:06:36 +00:00
|
|
|
|
2023-07-07 01:39:08 +00:00
|
|
|
_pairs, loras, inversions, _rest = parse_prompt(params)
|
2023-07-04 18:56:02 +00:00
|
|
|
for image, output in zip(images, outputs):
|
2023-07-10 06:02:30 +00:00
|
|
|
if full_res_inpaint:
|
|
|
|
if is_debug():
|
|
|
|
save_image(server, "adjusted-output.png", image)
|
|
|
|
mini_image = ImageOps.contain(image, (adj_mask_size, adj_mask_size))
|
|
|
|
image = original_source
|
|
|
|
image.paste(mini_image, box=adj_mask_border)
|
2023-07-04 18:56:02 +00:00
|
|
|
dest = save_image(
|
|
|
|
server,
|
|
|
|
output,
|
|
|
|
image,
|
|
|
|
params,
|
|
|
|
size,
|
|
|
|
upscale=upscale,
|
|
|
|
border=border,
|
|
|
|
inversions=inversions,
|
|
|
|
loras=loras,
|
|
|
|
)
|
2023-01-16 00:54:20 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# clean up
|
2023-01-22 04:28:13 +00:00
|
|
|
del image
|
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
|
2023-04-24 22:40:50 +00:00
|
|
|
show_system_toast(f"finished inpaint job: {dest}")
|
2023-02-05 13:53:26 +00:00
|
|
|
logger.info("finished inpaint job: %s", dest)
|
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-02-20 14:35:18 +00:00
|
|
|
outputs: List[str],
|
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
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# run and save
|
2023-07-15 23:54:54 +00:00
|
|
|
progress = worker.get_progress_callback()
|
2023-11-20 00:39:39 +00:00
|
|
|
images = chain.run(
|
|
|
|
worker, server, params, StageResult(images=[source]), callback=progress
|
|
|
|
)
|
2023-07-01 03:00:11 +00:00
|
|
|
|
2023-07-07 01:39:08 +00:00
|
|
|
_pairs, loras, inversions, _rest = parse_prompt(params)
|
2023-07-04 18:56:02 +00:00
|
|
|
for image, output in zip(images, outputs):
|
|
|
|
dest = save_image(
|
|
|
|
server,
|
|
|
|
output,
|
|
|
|
image,
|
|
|
|
params,
|
|
|
|
size,
|
|
|
|
upscale=upscale,
|
|
|
|
inversions=inversions,
|
|
|
|
loras=loras,
|
|
|
|
)
|
2023-04-15 01:29:44 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# clean up
|
2023-01-22 04:28:13 +00:00
|
|
|
del image
|
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
|
2023-04-24 22:40:50 +00:00
|
|
|
show_system_toast(f"finished upscale job: {dest}")
|
2023-02-05 13:53:26 +00:00
|
|
|
logger.info("finished upscale job: %s", dest)
|
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,
|
2023-02-20 14:35:18 +00:00
|
|
|
outputs: List[str],
|
2023-02-13 23:34:42 +00:00
|
|
|
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)
|
|
|
|
|
|
|
|
chain.stage(
|
|
|
|
BlendMaskStage(),
|
|
|
|
StageParams(tile_size=tile_size),
|
|
|
|
stage_source=sources[1],
|
|
|
|
stage_mask=mask,
|
|
|
|
)
|
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
|
|
|
)
|
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# run and save
|
2023-07-15 23:54:54 +00:00
|
|
|
progress = worker.get_progress_callback()
|
2023-11-20 00:39:39 +00:00
|
|
|
images = chain.run(
|
|
|
|
worker, server, params, StageResult(images=sources), callback=progress
|
|
|
|
)
|
2023-07-04 18:56:02 +00:00
|
|
|
|
|
|
|
for image, output in zip(images, outputs):
|
|
|
|
dest = save_image(server, output, image, params, size, upscale=upscale)
|
2023-02-13 23:34:42 +00:00
|
|
|
|
2023-06-30 04:06:36 +00:00
|
|
|
# clean up
|
2023-02-13 23:34:42 +00:00
|
|
|
del image
|
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
|
2023-04-24 22:40:50 +00:00
|
|
|
show_system_toast(f"finished blend job: {dest}")
|
2023-02-13 23:34:42 +00:00
|
|
|
logger.info("finished blend job: %s", dest)
|