mirror of
https://github.com/vladmandic/sdnext.git
synced 2026-01-27 15:02:48 +03:00
75 lines
3.6 KiB
Python
75 lines
3.6 KiB
Python
import time
|
|
from PIL import Image
|
|
from modules import shared, processing, images, sd_models, sd_vae
|
|
|
|
|
|
def get_tile(image: Image.Image, x: int, y: int, sx: int, sy: int) -> Image.Image:
|
|
return image.crop((
|
|
(x + 0) * image.width // sx,
|
|
(y + 0) * image.height // sy,
|
|
(x + 1) * image.width // sx,
|
|
(y + 1) * image.height // sy
|
|
))
|
|
|
|
|
|
def set_tile(image: Image.Image, x: int, y: int, tiled: Image.Image):
|
|
image.paste(tiled, (x * tiled.width, y * tiled.height))
|
|
return image
|
|
|
|
|
|
def run_tiling(p: processing.StableDiffusionProcessing, input_image: Image.Image) -> processing.Processed:
|
|
t0 = time.time()
|
|
# prepare images
|
|
sx, sy = p.control_tile.split('x')
|
|
sx = int(sx)
|
|
sy = int(sy)
|
|
vae_scale_factor = sd_vae.get_vae_scale_factor()
|
|
if sx <= 0 or sy <= 0:
|
|
raise ValueError('Control Tile: invalid tile size')
|
|
control_image = p.task_args.get('control_image', None) or p.task_args.get('image', None)
|
|
control_upscaled = None
|
|
if isinstance(control_image, list) and len(control_image) > 0:
|
|
w, h = vae_scale_factor * int(sx * control_image[0].width) // vae_scale_factor, vae_scale_factor * int(sy * control_image[0].height) // vae_scale_factor
|
|
control_upscaled = images.resize_image(resize_mode=1 if sx==sy else 5, im=control_image[0], width=w, height=h, context='add with forward')
|
|
init_image = p.override or input_image
|
|
init_upscaled = None
|
|
if init_image is not None:
|
|
w, h = vae_scale_factor * int(sx * init_image.width) // vae_scale_factor, vae_scale_factor * int(sy * init_image.height) // vae_scale_factor
|
|
init_upscaled = images.resize_image(resize_mode=1 if sx==sy else 5, im=init_image, width=w, height=h, context='add with forward')
|
|
t1 = time.time()
|
|
shared.log.debug(f'Control Tile: scale={sx}x{sy} resize={"fixed" if sx==sy else "context"} control={control_upscaled} init={init_upscaled} time={t1-t0:.3f}')
|
|
|
|
# stop processing from restoring pipeline on each iteration
|
|
orig_restore_pipeline = getattr(shared.sd_model, 'restore_pipeline', None)
|
|
shared.sd_model.restore_pipeline = None
|
|
|
|
# run tiling
|
|
for x in range(sx):
|
|
for y in range(sy):
|
|
shared.log.info(f'Control Tile: tile={x+1}-{sx}/{y+1}-{sy} target={control_upscaled}')
|
|
shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE)
|
|
p.init_images = None
|
|
p.task_args['control_mode'] = p.control_mode
|
|
p.task_args['strength'] = p.denoising_strength
|
|
if init_upscaled is not None:
|
|
p.task_args['image'] = [get_tile(init_upscaled, x, y, sx, sy)]
|
|
if control_upscaled is not None:
|
|
p.task_args['control_image'] = [get_tile(control_upscaled, x, y, sx, sy)]
|
|
processed: processing.Processed = processing.process_images(p) # run actual pipeline
|
|
if processed is None or len(processed.images) == 0:
|
|
continue
|
|
control_upscaled = set_tile(control_upscaled, x, y, processed.images[0])
|
|
|
|
# post-process
|
|
p.width = control_upscaled.width
|
|
p.height = control_upscaled.height
|
|
processed.images = [control_upscaled]
|
|
processed.info = processed.infotext(p, 0)
|
|
processed.infotexts = [processed.info]
|
|
shared.sd_model.restore_pipeline = orig_restore_pipeline
|
|
if hasattr(shared.sd_model, 'restore_pipeline') and shared.sd_model.restore_pipeline is not None:
|
|
shared.sd_model.restore_pipeline()
|
|
t2 = time.time()
|
|
shared.log.debug(f'Control Tile: image={control_upscaled} time={t2-t0:.3f}')
|
|
return processed
|