Spaces:
Running
Running
File size: 6,122 Bytes
613c9ab |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
from PIL import Image, ImageFilter
import torch
import math
from nodes import common_ksampler, VAEEncode, VAEDecode, VAEDecodeTiled
from utils import pil_to_tensor, tensor_to_pil, get_crop_region, expand_crop, crop_cond
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, uniform_tile_mode, tiled_decode):
# 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.uniform_tile_mode = uniform_tile_mode
self.tiled_decode = tiled_decode
self.vae_decoder = VAEDecode()
self.vae_encoder = VAEEncode()
self.vae_decoder_tiled = VAEDecodeTiled()
# 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)
if p.uniform_tile_mode:
# Expand the crop region to match the processing size ratio and then resize it to the processing size
x1, y1, x2, y2 = crop_region
crop_width = x2 - x1
crop_height = y2 - y1
crop_ratio = crop_width / crop_height
p_ratio = p.width / p.height
if crop_ratio > p_ratio:
target_width = crop_width
target_height = round(crop_width / p_ratio)
else:
target_width = round(crop_height * p_ratio)
target_height = crop_height
crop_region, _ = expand_crop(crop_region, image_mask.width, image_mask.height, target_width, target_height)
tile_size = p.width, p.height
else:
# Uses the minimal size that can fit the mask, minimizes tile size but may lead to image sizes that the model is not trained on
x1, y1, x2, y2 = crop_region
crop_width = x2 - x1
crop_height = y2 - y1
target_width = math.ceil(crop_width / 8) * 8
target_height = math.ceil(crop_height / 8) * 8
crop_region, tile_size = expand_crop(crop_region, image_mask.width,
image_mask.height, target_width, target_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]
# Assume the same size for all images in the batch
initial_tile_size = tiles[0].size
# Resize if necessary
for i, tile in enumerate(tiles):
if tile.size != tile_size:
tiles[i] = tile.resize(tile_size, Image.Resampling.LANCZOS)
# Crop conditioning
positive_cropped = crop_cond(p.positive, crop_region, p.init_size, init_image.size, tile_size)
negative_cropped = crop_cond(p.negative, crop_region, p.init_size, init_image.size, tile_size)
# Encode the image
batched_tiles = torch.cat([pil_to_tensor(tile) for tile in tiles], dim=0)
(latent,) = p.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
if not p.tiled_decode:
(decoded,) = p.vae_decoder.decode(p.vae, samples)
else:
print("[USDU] Using tiled decode")
(decoded,) = p.vae_decoder_tiled.decode(p.vae, samples, 512) # Default tile size is 512
# 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]
# 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
|