139 lines
5.0 KiB
Python
139 lines
5.0 KiB
Python
from PIL import Image, ImageFilter
|
|
import torch
|
|
from nodes import common_ksampler, VAEEncode, VAEDecode
|
|
from utils import pil_to_tensor, tensor_to_pil, get_crop_region, expand_crop, crop_cond, resize_and_pad_image
|
|
from modules import shared
|
|
|
|
if (not hasattr(Image, 'Resampling')): # For older versions of Pillow
|
|
Image.Resampling = Image
|
|
|
|
|
|
class StableDiffusionProcessing:
|
|
|
|
def __init__(self, init_img, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, upscale_by, force_uniform_tile_size):
|
|
# Variables used by the USDU script
|
|
self.init_images = [init_img]
|
|
self.image_mask = None
|
|
self.mask_blur = 0
|
|
self.inpaint_full_res_padding = 0
|
|
self.width = init_img.width
|
|
self.height = init_img.height
|
|
|
|
# ComfyUI Sampler inputs
|
|
self.model = model
|
|
self.positive = positive
|
|
self.negative = negative
|
|
self.vae = vae
|
|
self.seed = seed
|
|
self.steps = steps
|
|
self.cfg = cfg
|
|
self.sampler_name = sampler_name
|
|
self.scheduler = scheduler
|
|
self.denoise = denoise
|
|
|
|
# Variables used only by this script
|
|
self.init_size = init_img.width, init_img.height
|
|
self.upscale_by = upscale_by
|
|
self.force_uniform_tile_size = force_uniform_tile_size == "enable"
|
|
|
|
# Other required A1111 variables for the USDU script that is currently unused in this script
|
|
self.extra_generation_params = {}
|
|
|
|
|
|
class Processed:
|
|
|
|
def __init__(self, p: StableDiffusionProcessing, images: list, seed: int, info: str):
|
|
self.images = images
|
|
self.seed = seed
|
|
self.info = info
|
|
|
|
def infotext(self, p: StableDiffusionProcessing, index):
|
|
return None
|
|
|
|
|
|
def fix_seed(p: StableDiffusionProcessing):
|
|
pass
|
|
|
|
|
|
def process_images(p: StableDiffusionProcessing) -> Processed:
|
|
# Where the main image generation happens in A1111
|
|
|
|
# Setup
|
|
image_mask = p.image_mask.convert('L')
|
|
init_image = p.init_images[0]
|
|
|
|
# Locate the white region of the mask outlining the tile and add padding
|
|
crop_region = get_crop_region(image_mask, p.inpaint_full_res_padding)
|
|
crop_region, tile_size = expand_crop(crop_region, image_mask.width, image_mask.height)
|
|
|
|
# Blur the mask
|
|
if p.mask_blur > 0:
|
|
image_mask = image_mask.filter(ImageFilter.GaussianBlur(p.mask_blur))
|
|
|
|
# Crop the images to get the tiles that will be used for generation
|
|
tiles = [img.crop(crop_region) for img in shared.batch]
|
|
initial_tile_size = tiles[0].size
|
|
w_pad = 0
|
|
h_pad = 0
|
|
for i in range(len(tiles)):
|
|
if tiles[i].size != tile_size:
|
|
tiles[i] = tiles[i].resize(tile_size, Image.Resampling.LANCZOS)
|
|
|
|
if p.force_uniform_tile_size:
|
|
tiles[i], (w_pad, h_pad) = resize_and_pad_image(tiles[i], p.width, p.height, fill=True, blur=True)
|
|
|
|
# Crop conditioning
|
|
positive_cropped = crop_cond(p.positive, crop_region, p.init_size, init_image.size, tile_size, w_pad, h_pad)
|
|
negative_cropped = crop_cond(p.negative, crop_region, p.init_size, init_image.size, tile_size, w_pad, h_pad)
|
|
|
|
# Encode the image
|
|
vae_encoder = VAEEncode()
|
|
batched_tiles = torch.cat([pil_to_tensor(tile) for tile in tiles], dim=0)
|
|
(latent,) = vae_encoder.encode(p.vae, batched_tiles)
|
|
|
|
# Generate samples
|
|
(samples,) = common_ksampler(p.model, p.seed, p.steps, p.cfg, p.sampler_name,
|
|
p.scheduler, positive_cropped, negative_cropped, latent, denoise=p.denoise)
|
|
|
|
# Decode the sample
|
|
vae_decoder = VAEDecode()
|
|
(decoded,) = vae_decoder.decode(p.vae, samples)
|
|
|
|
# Convert the sample to a PIL image
|
|
tiles_sampled = [tensor_to_pil(decoded, i) for i in range(len(decoded))]
|
|
|
|
for i, tile_sampled in enumerate(tiles_sampled):
|
|
init_image = shared.batch[i]
|
|
|
|
if p.force_uniform_tile_size:
|
|
# Crop out the padding from the samples
|
|
tile_sampled = tile_sampled.crop((w_pad, h_pad, tile_sampled.width - w_pad, tile_sampled.height - h_pad))
|
|
# Resize the tile to the original size
|
|
tile_sampled = tile_sampled.resize(initial_tile_size, Image.Resampling.LANCZOS)
|
|
|
|
# Resize back to the original size
|
|
if tile_sampled.size != initial_tile_size:
|
|
tile_sampled = tile_sampled.resize(initial_tile_size, Image.Resampling.LANCZOS)
|
|
|
|
# Put the tile into position
|
|
image_tile_only = Image.new('RGBA', init_image.size)
|
|
image_tile_only.paste(tile_sampled, crop_region[:2])
|
|
|
|
# Add the mask as an alpha channel
|
|
# Must make a copy due to the possibility of an edge becoming black
|
|
temp = image_tile_only.copy()
|
|
temp.putalpha(image_mask)
|
|
image_tile_only.paste(temp, image_tile_only)
|
|
|
|
# Add back the tile to the initial image according to the mask in the alpha channel
|
|
result = init_image.convert('RGBA')
|
|
result.alpha_composite(image_tile_only)
|
|
|
|
# Convert back to RGB
|
|
result = result.convert('RGB')
|
|
|
|
shared.batch[i] = result
|
|
|
|
processed = Processed(p, [shared.batch[0]], p.seed, None)
|
|
return processed
|