diff --git a/modules/processing.py b/modules/processing.py index d888ffa..b9f5eb5 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -1,11 +1,33 @@ # Patched classes to adapt from A111 webui for ComfyUI +from nodes import common_ksampler, VAEEncodeTiled, VAEDecodeTiled, ConditioningSetMask +from utils import pil_to_tensor, tensor_to_pil +import modules.shared as shared +import numpy as np +import torch +from PIL import Image + class StableDiffusionProcessing: - seed = 0 - extra_generation_params = {} - def __init__(self, init_img): + def __init__(self, init_img, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise): + # Variables used by the upscaler script self.init_images = [init_img] + self.image_mask = None + + # 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 + + # Other required A1111 variables for the upscaler script that is currently unused in this script + self.extra_generation_params = {} class Processed: @@ -22,9 +44,47 @@ class Processed: def fix_seed(p: StableDiffusionProcessing): pass - def process_images(p: StableDiffusionProcessing) -> Processed: # Where the main image generation happens in A1111 - # Return original images for now - processed = Processed(p, p.init_images, p.seed, None) + + # Convert the PIL images to a torch tensor + init_images = p.init_images + image_tensor = pil_to_tensor(init_images[0]) + + # Encode the image + vae_encoder = VAEEncodeTiled() + (encoded,) = vae_encoder.encode(p.vae, image_tensor) + print(encoded["samples"].shape) + + # Convert the black and white mask to a torch tensor + mask_pil = p.image_mask + mask_pil_mono = mask_pil.convert("L") + mask = np.array(mask_pil_mono).astype(np.float32) / 255.0 + mask = torch.from_numpy(mask) + + # Add the mask to the conditioning + conditioning_set_mask = ConditioningSetMask() + (masked_positive,) = conditioning_set_mask.append(p.positive, mask, "mask bounds", 1) + (masked_negative,) = conditioning_set_mask.append(p.negative, mask, "mask bounds", 1) + + # Generate samples + (samples,) = common_ksampler(p.model, p.seed, p.steps, p.cfg, p.sampler_name, + p.scheduler, masked_positive, masked_negative, encoded, denoise=p.denoise) + + # Decode the sample + vae_decoder = VAEDecodeTiled() + (decoded,) = vae_decoder.decode(p.vae, samples) + + # Convert the sample to a PIL image + image = tensor_to_pil(decoded) + + # Because ComfyUI noises the masked parts of the image as well, the image must be assembled elsewhere + if shared.tiled_image is None: + shared.tiled_image = image + else: + # Add the tile to the tiled image using the mask + shared.tiled_image = Image.composite(image, shared.tiled_image, mask_pil_mono) + + # Return the original image instead of the generated image because the masked parts of the image are noised + processed = Processed(p, init_images, p.seed, None) return processed diff --git a/modules/shared.py b/modules/shared.py index b6e16e8..6b4f123 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -19,3 +19,6 @@ state = State() sd_upscalers = [None] # The upscaler usable by ComfyUI nodes actual_upscaler = None + +# The resulting image +tiled_image = None \ No newline at end of file diff --git a/nodes.py b/nodes.py index 2c89311..6c50d96 100644 --- a/nodes.py +++ b/nodes.py @@ -14,7 +14,11 @@ sys.path.insert(0, os.path.join( MAX_RESOLUTION = 8192 # The modes avaiable for Ultimate SD Upscale -MODES = ["Linear", "Chess", "None"] +MODES = { + "Linear": ult.USDUMode.LINEAR, + "Chess": ult.USDUMode.CHESS, + "None": ult.USDUMode.NONE, +} class UltimateSDUpscale: @@ -24,9 +28,10 @@ class UltimateSDUpscale: "required": { "image": ("IMAGE",), # Sampling Params - # "model": ("MODEL",), - # "positive": ("CONDITIONING",), - # "negative": ("CONDITIONING",), + "model": ("MODEL",), + "positive": ("CONDITIONING",), + "negative": ("CONDITIONING",), + "vae": ("VAE",), "upscale_by": ("FLOAT", {"default": 2, "min": 0.05, "max": 4, "step": 0.05}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), "steps": ("INT", {"default": 20, "min": 1, "max": 10000, "step": 1}), @@ -36,7 +41,7 @@ class UltimateSDUpscale: "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), # Upscale Params "upscale_model": ("UPSCALE_MODEL",), - "mode_type": (MODES,), + "mode_type": (list(MODES.keys()),), "tile_width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}), "tile_height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}), "mask_blur": ("INT", {"default": 8, "min": 0, "max": 64, "step": 1}), @@ -48,9 +53,9 @@ class UltimateSDUpscale: FUNCTION = "upscale" CATEGORY = "image/upscaling" - def upscale(self, image, # model, positive, negative, - upscale_by, seed, steps, cfg, sampler_name, scheduler, denoise, - upscale_model, mode_type, tile_width, tile_height, mask_blur, tile_padding): + def upscale(self, image, model, positive, negative, vae, upscale_by, seed, + steps, cfg, sampler_name, scheduler, denoise, upscale_model, + mode_type, tile_width, tile_height, mask_blur, tile_padding): # # Set up A1111 patches # @@ -60,9 +65,14 @@ class UltimateSDUpscale: shared.sd_upscalers[0] = UpscalerData() # Where the actual upscaler is stored, will be used when the script upscales using the Upscaler in UpscalerData shared.actual_upscaler = upscale_model + # Reset the resulting image + shared.tiled_image = None # Processing - sdprocessing = StableDiffusionProcessing(tensor_to_pil(image)) + sdprocessing = StableDiffusionProcessing(tensor_to_pil(image), + model, positive, negative, vae, + seed, steps, cfg, sampler_name, + scheduler, denoise) # # Running the script @@ -70,12 +80,12 @@ class UltimateSDUpscale: script = ult.Script() processed = script.run(p=sdprocessing, _=None, tile_width=tile_width, tile_height=tile_height, mask_blur=mask_blur, padding=tile_padding, seams_fix_width=None, seams_fix_denoise=None, seams_fix_padding=None, - upscaler_index=0, save_upscaled_image=False, redraw_mode=ult.USDUMode.LINEAR, save_seams_fix_image=False, + upscaler_index=0, save_upscaled_image=False, redraw_mode=MODES[mode_type], save_seams_fix_image=False, seams_fix_mask_blur=None, seams_fix_type=ult.USDUSFMode.NONE, target_size_type=2, custom_width=None, custom_height=None, custom_scale=upscale_by) - # Return the upscaled image - upscaled_image = pil_to_tensor(processed.images[0]) + # Return the resulting image + upscaled_image = pil_to_tensor(shared.tiled_image) return (upscaled_image,) diff --git a/utils.py b/utils.py index 0a1fe90..d76ccfd 100644 --- a/utils.py +++ b/utils.py @@ -3,11 +3,13 @@ import PIL.Image as Image import torch def tensor_to_pil(img_tensor): + # Takes a batch of 1 rgb image and returns an RGB PIL image i = 255. * img_tensor.cpu().numpy() img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8).squeeze()) return img def pil_to_tensor(img): + # Takes a 3 channel PIL image and returns a tensor of shape [1, height, width, 3] image = img.convert("RGB") image = np.array(image).astype(np.float32) / 255.0 image = torch.from_numpy(image)[None,]