代码上传
This commit is contained in:
m18328015652@163.com
@@ -0,0 +1,2 @@
|
||||
def torch_gc():
|
||||
pass
|
||||
@@ -0,0 +1,8 @@
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def flatten(img, bgcolor):
|
||||
# Replace transparency with bgcolor
|
||||
if img.mode in ("RGB"):
|
||||
return img
|
||||
return Image.alpha_composite(Image.new("RGBA", img.size, bgcolor), img).convert("RGB")
|
||||
@@ -0,0 +1,163 @@
|
||||
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 . 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
|
||||
@@ -0,0 +1,473 @@
|
||||
import math
|
||||
from PIL import Image, ImageDraw, ImageOps
|
||||
from .processing import StableDiffusionProcessing
|
||||
from .processing import Processed
|
||||
from enum import Enum
|
||||
from . import shared
|
||||
from . import images
|
||||
from . import devices
|
||||
from . import processing
|
||||
|
||||
class USDUMode(Enum):
|
||||
LINEAR = 0
|
||||
CHESS = 1
|
||||
NONE = 2
|
||||
|
||||
class USDUSFMode(Enum):
|
||||
NONE = 0
|
||||
BAND_PASS = 1
|
||||
HALF_TILE = 2
|
||||
HALF_TILE_PLUS_INTERSECTIONS = 3
|
||||
|
||||
class USDUpscaler():
|
||||
|
||||
def __init__(self, p, image, upscaler_index:int, save_redraw, save_seams_fix, tile_width, tile_height) -> None:
|
||||
self.p:StableDiffusionProcessing = p
|
||||
self.image:Image = image
|
||||
self.scale_factor = math.ceil(max(p.width, p.height) / max(image.width, image.height))
|
||||
self.upscaler = shared.sd_upscalers[upscaler_index]
|
||||
self.redraw = USDURedraw()
|
||||
self.redraw.save = save_redraw
|
||||
self.redraw.tile_width = tile_width if tile_width > 0 else tile_height
|
||||
self.redraw.tile_height = tile_height if tile_height > 0 else tile_width
|
||||
self.seams_fix = USDUSeamsFix()
|
||||
self.seams_fix.save = save_seams_fix
|
||||
self.seams_fix.tile_width = tile_width if tile_width > 0 else tile_height
|
||||
self.seams_fix.tile_height = tile_height if tile_height > 0 else tile_width
|
||||
self.initial_info = None
|
||||
self.rows = math.ceil(self.p.height / 4 / self.redraw.tile_height)
|
||||
self.cols = math.ceil(self.p.width / self.redraw.tile_width)
|
||||
|
||||
def get_factor(self, num):
|
||||
# Its just return, don't need elif
|
||||
if num == 1:
|
||||
return 2
|
||||
if num % 4 == 0:
|
||||
return 4
|
||||
if num % 3 == 0:
|
||||
return 3
|
||||
if num % 2 == 0:
|
||||
return 2
|
||||
return 0
|
||||
|
||||
def get_factors(self):
|
||||
scales = []
|
||||
current_scale = 1
|
||||
current_scale_factor = self.get_factor(self.scale_factor)
|
||||
while current_scale_factor == 0:
|
||||
self.scale_factor += 1
|
||||
current_scale_factor = self.get_factor(self.scale_factor)
|
||||
while current_scale < self.scale_factor:
|
||||
current_scale_factor = self.get_factor(self.scale_factor // current_scale)
|
||||
scales.append(current_scale_factor)
|
||||
current_scale = current_scale * current_scale_factor
|
||||
if current_scale_factor == 0:
|
||||
break
|
||||
self.scales = enumerate(scales)
|
||||
|
||||
def upscale(self):
|
||||
# Log info
|
||||
print(f"Canva size: {self.p.width}x{self.p.height}")
|
||||
print(f"Image size: {self.image.width}x{self.image.height}")
|
||||
print(f"Scale factor: {self.scale_factor}")
|
||||
# Check upscaler is not empty
|
||||
if self.upscaler.name == "None":
|
||||
self.image = self.image.resize((self.p.width, self.p.height), resample=Image.LANCZOS)
|
||||
return
|
||||
# Get list with scale factors
|
||||
self.get_factors()
|
||||
# Upscaling image over all factors
|
||||
for index, value in self.scales:
|
||||
print(f"Upscaling iteration {index+1} with scale factor {value}")
|
||||
self.image = self.upscaler.scaler.upscale(self.image, value, self.upscaler.data_path)
|
||||
# Resize image to set values
|
||||
self.image = self.image.resize((self.p.width, self.p.height), resample=Image.LANCZOS)
|
||||
shared.batch = [self.image] + [img.resize((self.p.width, self.p.height), resample=Image.LANCZOS) for img in shared.batch[1:]]
|
||||
|
||||
|
||||
def setup_redraw(self, redraw_mode, padding, mask_blur):
|
||||
self.redraw.mode = USDUMode(redraw_mode)
|
||||
self.redraw.enabled = self.redraw.mode != USDUMode.NONE
|
||||
self.redraw.padding = padding
|
||||
self.p.mask_blur = mask_blur
|
||||
|
||||
|
||||
def setup_seams_fix(self, padding, denoise, mask_blur, width, mode):
|
||||
self.seams_fix.padding = padding
|
||||
self.seams_fix.denoise = denoise
|
||||
self.seams_fix.mask_blur = mask_blur
|
||||
self.seams_fix.width = width
|
||||
self.seams_fix.mode = USDUSFMode(mode)
|
||||
self.seams_fix.enabled = self.seams_fix.mode != USDUSFMode.NONE
|
||||
|
||||
|
||||
def save_image(self):
|
||||
if type(self.p.prompt) != list:
|
||||
images.save_image(self.image, self.p.outpath_samples, "", self.p.seed, self.p.prompt, shared.opts.samples_format, info=self.initial_info, p=self.p)
|
||||
else:
|
||||
images.save_image(self.image, self.p.outpath_samples, "", self.p.seed, self.p.prompt[0], shared.opts.samples_format, info=self.initial_info, p=self.p)
|
||||
|
||||
def calc_jobs_count(self):
|
||||
redraw_job_count = (self.rows * self.cols) if self.redraw.enabled else 0
|
||||
seams_job_count = 0
|
||||
if self.seams_fix.mode == USDUSFMode.BAND_PASS:
|
||||
seams_job_count = self.rows + self.cols - 2
|
||||
elif self.seams_fix.mode == USDUSFMode.HALF_TILE:
|
||||
seams_job_count = self.rows * (self.cols - 1) + (self.rows - 1) * self.cols
|
||||
elif self.seams_fix.mode == USDUSFMode.HALF_TILE_PLUS_INTERSECTIONS:
|
||||
seams_job_count = self.rows * (self.cols - 1) + (self.rows - 1) * self.cols + (self.rows - 1) * (self.cols - 1)
|
||||
|
||||
shared.state.job_count = redraw_job_count + seams_job_count
|
||||
|
||||
def print_info(self):
|
||||
print(f"Tile size: {self.redraw.tile_width}x{self.redraw.tile_height}")
|
||||
print(f"Tiles amount: {self.rows * self.cols}")
|
||||
print(f"Grid: {self.rows}x{self.cols}")
|
||||
print(f"Redraw enabled: {self.redraw.enabled}")
|
||||
print(f"Seams fix mode: {self.seams_fix.mode.name}")
|
||||
|
||||
def add_extra_info(self):
|
||||
self.p.extra_generation_params["Ultimate SD upscale upscaler"] = self.upscaler.name
|
||||
self.p.extra_generation_params["Ultimate SD upscale tile_width"] = self.redraw.tile_width
|
||||
self.p.extra_generation_params["Ultimate SD upscale tile_height"] = self.redraw.tile_height
|
||||
self.p.extra_generation_params["Ultimate SD upscale mask_blur"] = self.p.mask_blur
|
||||
self.p.extra_generation_params["Ultimate SD upscale padding"] = self.redraw.padding
|
||||
|
||||
def process(self):
|
||||
shared.state.begin()
|
||||
self.calc_jobs_count()
|
||||
self.result_images = []
|
||||
if self.redraw.enabled:
|
||||
self.image = self.redraw.start(self.p, self.image, self.rows, self.cols)
|
||||
self.initial_info = self.redraw.initial_info
|
||||
self.result_images.append(self.image)
|
||||
if self.redraw.save:
|
||||
self.save_image()
|
||||
|
||||
if self.seams_fix.enabled:
|
||||
self.image = self.seams_fix.start(self.p, self.image, self.rows, self.cols)
|
||||
self.initial_info = self.seams_fix.initial_info
|
||||
self.result_images.append(self.image)
|
||||
if self.seams_fix.save:
|
||||
self.save_image()
|
||||
shared.state.end()
|
||||
|
||||
class USDURedraw():
|
||||
|
||||
def init_draw(self, p, width, height):
|
||||
p.inpaint_full_res = True
|
||||
p.inpaint_full_res_padding = self.padding
|
||||
p.width = math.ceil((self.tile_width+self.padding) / 8) * 8
|
||||
p.height = math.ceil((self.tile_height+self.padding) / 8) * 8
|
||||
mask = Image.new("L", (width, height), "black")
|
||||
draw = ImageDraw.Draw(mask)
|
||||
return mask, draw
|
||||
|
||||
def calc_rectangle(self, xi, yi):
|
||||
x1 = xi * self.tile_width
|
||||
y1 = yi * self.tile_height
|
||||
x2 = xi * self.tile_width + self.tile_width
|
||||
y2 = yi * self.tile_height + self.tile_height
|
||||
|
||||
return x1, y1, x2, y2
|
||||
|
||||
def linear_process(self, p, image, rows, cols):
|
||||
mask, draw = self.init_draw(p, image.width, image.height)
|
||||
for yi in range(rows):
|
||||
for xi in range(cols):
|
||||
if shared.state.interrupted:
|
||||
break
|
||||
draw.rectangle(self.calc_rectangle(xi, yi), fill="white")
|
||||
p.init_images = [image]
|
||||
p.image_mask = mask
|
||||
processed = processing.process_images(p)
|
||||
draw.rectangle(self.calc_rectangle(xi, yi), fill="black")
|
||||
if (len(processed.images) > 0):
|
||||
image = processed.images[0]
|
||||
|
||||
p.width = image.width
|
||||
p.height = image.height
|
||||
self.initial_info = processed.infotext(p, 0)
|
||||
|
||||
return image
|
||||
|
||||
def chess_process(self, p, image, rows, cols):
|
||||
mask, draw = self.init_draw(p, image.width, image.height)
|
||||
tiles = []
|
||||
# calc tiles colors
|
||||
for yi in range(rows):
|
||||
for xi in range(cols):
|
||||
if shared.state.interrupted:
|
||||
break
|
||||
if xi == 0:
|
||||
tiles.append([])
|
||||
color = xi % 2 == 0
|
||||
if yi > 0 and yi % 2 != 0:
|
||||
color = not color
|
||||
tiles[yi].append(color)
|
||||
|
||||
for yi in range(len(tiles)):
|
||||
for xi in range(len(tiles[yi])):
|
||||
if shared.state.interrupted:
|
||||
break
|
||||
if not tiles[yi][xi]:
|
||||
tiles[yi][xi] = not tiles[yi][xi]
|
||||
continue
|
||||
tiles[yi][xi] = not tiles[yi][xi]
|
||||
draw.rectangle(self.calc_rectangle(xi, yi), fill="white")
|
||||
p.init_images = [image]
|
||||
p.image_mask = mask
|
||||
processed = processing.process_images(p)
|
||||
draw.rectangle(self.calc_rectangle(xi, yi), fill="black")
|
||||
if (len(processed.images) > 0):
|
||||
image = processed.images[0]
|
||||
|
||||
for yi in range(len(tiles)):
|
||||
for xi in range(len(tiles[yi])):
|
||||
if shared.state.interrupted:
|
||||
break
|
||||
if not tiles[yi][xi]:
|
||||
continue
|
||||
draw.rectangle(self.calc_rectangle(xi, yi), fill="white")
|
||||
p.init_images = [image]
|
||||
p.image_mask = mask
|
||||
processed = processing.process_images(p)
|
||||
draw.rectangle(self.calc_rectangle(xi, yi), fill="black")
|
||||
if (len(processed.images) > 0):
|
||||
image = processed.images[0]
|
||||
|
||||
p.width = image.width
|
||||
p.height = image.height
|
||||
self.initial_info = processed.infotext(p, 0)
|
||||
|
||||
return image
|
||||
|
||||
def start(self, p, image, rows, cols):
|
||||
self.initial_info = None
|
||||
if self.mode == USDUMode.LINEAR:
|
||||
return self.linear_process(p, image, rows, cols)
|
||||
if self.mode == USDUMode.CHESS:
|
||||
return self.chess_process(p, image, rows, cols)
|
||||
|
||||
class USDUSeamsFix():
|
||||
|
||||
def init_draw(self, p):
|
||||
self.initial_info = None
|
||||
p.width = math.ceil((self.tile_width+self.padding) / 8) * 8
|
||||
p.height = math.ceil((self.tile_height+self.padding) / 8) * 8
|
||||
|
||||
def half_tile_process(self, p, image, rows, cols):
|
||||
|
||||
self.init_draw(p)
|
||||
processed = None
|
||||
|
||||
gradient = Image.linear_gradient("L")
|
||||
row_gradient = Image.new("L", (self.tile_width, self.tile_height), "black")
|
||||
row_gradient.paste(gradient.resize(
|
||||
(self.tile_width, self.tile_height//2), resample=Image.BICUBIC), (0, 0))
|
||||
row_gradient.paste(gradient.rotate(180).resize(
|
||||
(self.tile_width, self.tile_height//2), resample=Image.BICUBIC),
|
||||
(0, self.tile_height//2))
|
||||
col_gradient = Image.new("L", (self.tile_width, self.tile_height), "black")
|
||||
col_gradient.paste(gradient.rotate(90).resize(
|
||||
(self.tile_width//2, self.tile_height), resample=Image.BICUBIC), (0, 0))
|
||||
col_gradient.paste(gradient.rotate(270).resize(
|
||||
(self.tile_width//2, self.tile_height), resample=Image.BICUBIC), (self.tile_width//2, 0))
|
||||
|
||||
p.denoising_strength = self.denoise
|
||||
p.mask_blur = self.mask_blur
|
||||
|
||||
for yi in range(rows-1):
|
||||
for xi in range(cols):
|
||||
if shared.state.interrupted:
|
||||
break
|
||||
p.width = self.tile_width
|
||||
p.height = self.tile_height
|
||||
p.inpaint_full_res = True
|
||||
p.inpaint_full_res_padding = self.padding
|
||||
mask = Image.new("L", (image.width, image.height), "black")
|
||||
mask.paste(row_gradient, (xi*self.tile_width, yi*self.tile_height + self.tile_height//2))
|
||||
|
||||
p.init_images = [image]
|
||||
p.image_mask = mask
|
||||
processed = processing.process_images(p)
|
||||
if (len(processed.images) > 0):
|
||||
image = processed.images[0]
|
||||
|
||||
for yi in range(rows):
|
||||
for xi in range(cols-1):
|
||||
if shared.state.interrupted:
|
||||
break
|
||||
p.width = self.tile_width
|
||||
p.height = self.tile_height
|
||||
p.inpaint_full_res = True
|
||||
p.inpaint_full_res_padding = self.padding
|
||||
mask = Image.new("L", (image.width, image.height), "black")
|
||||
mask.paste(col_gradient, (xi*self.tile_width+self.tile_width//2, yi*self.tile_height))
|
||||
|
||||
p.init_images = [image]
|
||||
p.image_mask = mask
|
||||
processed = processing.process_images(p)
|
||||
if (len(processed.images) > 0):
|
||||
image = processed.images[0]
|
||||
|
||||
p.width = image.width
|
||||
p.height = image.height
|
||||
if processed is not None:
|
||||
self.initial_info = processed.infotext(p, 0)
|
||||
|
||||
return image
|
||||
|
||||
def half_tile_process_corners(self, p, image, rows, cols):
|
||||
fixed_image = self.half_tile_process(p, image, rows, cols)
|
||||
processed = None
|
||||
self.init_draw(p)
|
||||
gradient = Image.radial_gradient("L").resize(
|
||||
(self.tile_width, self.tile_height), resample=Image.BICUBIC)
|
||||
gradient = ImageOps.invert(gradient)
|
||||
p.denoising_strength = self.denoise
|
||||
#p.mask_blur = 0
|
||||
p.mask_blur = self.mask_blur
|
||||
|
||||
for yi in range(rows-1):
|
||||
for xi in range(cols-1):
|
||||
if shared.state.interrupted:
|
||||
break
|
||||
p.width = self.tile_width
|
||||
p.height = self.tile_height
|
||||
p.inpaint_full_res = True
|
||||
p.inpaint_full_res_padding = 0
|
||||
mask = Image.new("L", (fixed_image.width, fixed_image.height), "black")
|
||||
mask.paste(gradient, (xi*self.tile_width + self.tile_width//2,
|
||||
yi*self.tile_height + self.tile_height//2))
|
||||
|
||||
p.init_images = [fixed_image]
|
||||
p.image_mask = mask
|
||||
processed = processing.process_images(p)
|
||||
if (len(processed.images) > 0):
|
||||
fixed_image = processed.images[0]
|
||||
|
||||
p.width = fixed_image.width
|
||||
p.height = fixed_image.height
|
||||
if processed is not None:
|
||||
self.initial_info = processed.infotext(p, 0)
|
||||
|
||||
return fixed_image
|
||||
|
||||
def band_pass_process(self, p, image, cols, rows):
|
||||
|
||||
self.init_draw(p)
|
||||
processed = None
|
||||
|
||||
p.denoising_strength = self.denoise
|
||||
p.mask_blur = 0
|
||||
|
||||
gradient = Image.linear_gradient("L")
|
||||
mirror_gradient = Image.new("L", (256, 256), "black")
|
||||
mirror_gradient.paste(gradient.resize((256, 128), resample=Image.BICUBIC), (0, 0))
|
||||
mirror_gradient.paste(gradient.rotate(180).resize((256, 128), resample=Image.BICUBIC), (0, 128))
|
||||
|
||||
row_gradient = mirror_gradient.resize((image.width, self.width), resample=Image.BICUBIC)
|
||||
col_gradient = mirror_gradient.rotate(90).resize((self.width, image.height), resample=Image.BICUBIC)
|
||||
|
||||
for xi in range(1, rows):
|
||||
if shared.state.interrupted:
|
||||
break
|
||||
p.width = self.width + self.padding * 2
|
||||
p.height = image.height
|
||||
p.inpaint_full_res = True
|
||||
p.inpaint_full_res_padding = self.padding
|
||||
mask = Image.new("L", (image.width, image.height), "black")
|
||||
mask.paste(col_gradient, (xi * self.tile_width - self.width // 2, 0))
|
||||
|
||||
p.init_images = [image]
|
||||
p.image_mask = mask
|
||||
processed = processing.process_images(p)
|
||||
if (len(processed.images) > 0):
|
||||
image = processed.images[0]
|
||||
for yi in range(1, cols):
|
||||
if shared.state.interrupted:
|
||||
break
|
||||
p.width = image.width
|
||||
p.height = self.width + self.padding * 2
|
||||
p.inpaint_full_res = True
|
||||
p.inpaint_full_res_padding = self.padding
|
||||
mask = Image.new("L", (image.width, image.height), "black")
|
||||
mask.paste(row_gradient, (0, yi * self.tile_height - self.width // 2))
|
||||
|
||||
p.init_images = [image]
|
||||
p.image_mask = mask
|
||||
processed = processing.process_images(p)
|
||||
if (len(processed.images) > 0):
|
||||
image = processed.images[0]
|
||||
|
||||
p.width = image.width
|
||||
p.height = image.height
|
||||
if processed is not None:
|
||||
self.initial_info = processed.infotext(p, 0)
|
||||
|
||||
return image
|
||||
|
||||
def start(self, p, image, rows, cols):
|
||||
if USDUSFMode(self.mode) == USDUSFMode.BAND_PASS:
|
||||
return self.band_pass_process(p, image, rows, cols)
|
||||
elif USDUSFMode(self.mode) == USDUSFMode.HALF_TILE:
|
||||
return self.half_tile_process(p, image, rows, cols)
|
||||
elif USDUSFMode(self.mode) == USDUSFMode.HALF_TILE_PLUS_INTERSECTIONS:
|
||||
return self.half_tile_process_corners(p, image, rows, cols)
|
||||
else:
|
||||
return image
|
||||
|
||||
class Script():
|
||||
def title(self):
|
||||
return "Ultimate SD upscale"
|
||||
|
||||
def show(self, is_img2img):
|
||||
return is_img2img
|
||||
|
||||
def run(self, p, _, tile_width, tile_height, mask_blur, padding, seams_fix_width, seams_fix_denoise, seams_fix_padding,
|
||||
upscaler_index, save_upscaled_image, redraw_mode, save_seams_fix_image, seams_fix_mask_blur,
|
||||
seams_fix_type, target_size_type, custom_width, custom_height, custom_scale):
|
||||
|
||||
# Init
|
||||
processing.fix_seed(p)
|
||||
devices.torch_gc()
|
||||
|
||||
p.do_not_save_grid = True
|
||||
p.do_not_save_samples = True
|
||||
p.inpaint_full_res = False
|
||||
|
||||
p.inpainting_fill = 1
|
||||
p.n_iter = 1
|
||||
p.batch_size = 1
|
||||
|
||||
seed = p.seed
|
||||
|
||||
# Init image
|
||||
init_img = p.init_images[0]
|
||||
if init_img == None:
|
||||
return Processed(p, [], seed, "Empty image")
|
||||
init_img = images.flatten(init_img, shared.opts.img2img_background_color)
|
||||
|
||||
#override size
|
||||
if target_size_type == 1:
|
||||
p.width = custom_width
|
||||
p.height = custom_height
|
||||
if target_size_type == 2:
|
||||
p.width = math.ceil((init_img.width * custom_scale) / 8) * 8
|
||||
p.height = math.ceil((init_img.height * custom_scale) / 8) * 8
|
||||
|
||||
# Upscaling
|
||||
upscaler = USDUpscaler(p, init_img, upscaler_index, save_upscaled_image, save_seams_fix_image, tile_width, tile_height)
|
||||
upscaler.upscale()
|
||||
|
||||
# Drawing
|
||||
upscaler.setup_redraw(redraw_mode, padding, mask_blur)
|
||||
upscaler.setup_seams_fix(seams_fix_padding, seams_fix_denoise, seams_fix_mask_blur, seams_fix_width, seams_fix_type)
|
||||
upscaler.print_info()
|
||||
upscaler.add_extra_info()
|
||||
upscaler.process()
|
||||
result_images = upscaler.result_images
|
||||
|
||||
return Processed(p, result_images, seed, upscaler.initial_info if upscaler.initial_info is not None else "")
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
class Options:
|
||||
img2img_background_color = "#ffffff" # Set to white for now
|
||||
|
||||
|
||||
class State:
|
||||
interrupted = False
|
||||
|
||||
def begin(self):
|
||||
pass
|
||||
|
||||
def end(self):
|
||||
pass
|
||||
|
||||
|
||||
opts = Options()
|
||||
state = State()
|
||||
|
||||
# Will only ever hold 1 upscaler
|
||||
sd_upscalers = [None]
|
||||
# The upscaler usable by ComfyUI nodes
|
||||
actual_upscaler = None
|
||||
|
||||
# Batch of images to upscale
|
||||
batch = None
|
||||
@@ -0,0 +1,30 @@
|
||||
from PIL import Image
|
||||
from ..utils import tensor_to_pil, pil_to_tensor
|
||||
from comfy_extras.nodes_upscale_model import ImageUpscaleWithModel
|
||||
from ..modules import shared
|
||||
|
||||
if (not hasattr(Image, 'Resampling')): # For older versions of Pillow
|
||||
Image.Resampling = Image
|
||||
|
||||
|
||||
class Upscaler:
|
||||
|
||||
def _upscale(self, img: Image, scale):
|
||||
if (shared.actual_upscaler is None):
|
||||
return img.resize((img.width * scale, img.height * scale), Image.Resampling.LANCZOS)
|
||||
tensor = pil_to_tensor(img)
|
||||
image_upscale_node = ImageUpscaleWithModel()
|
||||
(upscaled,) = image_upscale_node.upscale(shared.actual_upscaler, tensor)
|
||||
return tensor_to_pil(upscaled)
|
||||
|
||||
def upscale(self, img: Image, scale, selected_model: str = None):
|
||||
shared.batch = [self._upscale(img, scale) for img in shared.batch]
|
||||
return shared.batch[0]
|
||||
|
||||
|
||||
class UpscalerData:
|
||||
name = ""
|
||||
data_path = ""
|
||||
|
||||
def __init__(self):
|
||||
self.scaler = Upscaler()
|
||||
Reference in New Issue
Block a user