From b30d4a84bde8f8df9b35bd9457cc9adeb5d2e247 Mon Sep 17 00:00:00 2001 From: matt3o Date: Mon, 4 Mar 2024 14:20:53 +0100 Subject: [PATCH] fix variations with noise --- essentials.py | 65 ++++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 52 insertions(+), 13 deletions(-) diff --git a/essentials.py b/essentials.py index 78cb352..6652f35 100644 --- a/essentials.py +++ b/essentials.py @@ -6,6 +6,7 @@ import random import os import operator as op import numpy as np +import scipy from PIL import Image, ImageDraw, ImageFont, ImageColor, ImageFilter import io @@ -17,6 +18,7 @@ from nodes import MAX_RESOLUTION, SaveImage, common_ksampler import folder_paths import comfy.utils import comfy.samplers +import comfy.sample STOCHASTIC_SAMPLERS = ["euler_ancestral", "dpm_2_ancestral", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm"] @@ -1088,6 +1090,32 @@ def slerp(val, low, high): return res.reshape(dims) +def prepare_mask(mask, shape): + mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[2], shape[3]), mode="bilinear") + mask = mask.expand((-1,shape[1],-1,-1)) + if mask.shape[0] < shape[0]: + mask = mask.repeat((shape[0] -1) // mask.shape[0] + 1, 1, 1, 1)[:shape[0]] + return mask + +def expand_mask(mask, expand, tapered_corners): + c = 0 if tapered_corners else 1 + kernel = np.array([[c, 1, c], + [1, 1, 1], + [c, 1, c]]) + mask = mask.reshape((-1, mask.shape[-2], mask.shape[-1])) + out = [] + for m in mask: + output = m.numpy() + for _ in range(abs(expand)): + if expand < 0: + output = scipy.ndimage.grey_erosion(output, footprint=kernel) + else: + output = scipy.ndimage.grey_dilation(output, footprint=kernel) + output = torch.from_numpy(output) + out.append(output) + + return torch.stack(out, dim=0) + class KSamplerVariationsWithNoise: @classmethod def INPUT_TYPES(s): @@ -1101,31 +1129,42 @@ class KSamplerVariationsWithNoise: "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ), "positive": ("CONDITIONING", ), "negative": ("CONDITIONING", ), - "variation_strength": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step":0.01, "round": 0.01}), + "variation_strength": ("FLOAT", {"default": 0.17, "min": 0.0, "max": 1.0, "step":0.01, "round": 0.01}), #"start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}), #"end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}), #"return_with_leftover_noise": (["disable", "enable"], ), - "variation_seed": ("INT:seed", {"default": random.randint(0, 0xffffffffffffffff), "min": 0, "max": 0xffffffffffffffff}), + "variation_seed": ("INT:seed", {"default": 12345, "min": 0, "max": 0xffffffffffffffff}), + "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step":0.01, "round": 0.01}), }} RETURN_TYPES = ("LATENT",) FUNCTION = "execute" CATEGORY = "essentials" - def execute(self, model, latent_image, main_seed, steps, cfg, sampler_name, scheduler, positive, negative, variation_strength, variation_seed): - generator = torch.manual_seed(main_seed) + def execute(self, model, latent_image, main_seed, steps, cfg, sampler_name, scheduler, positive, negative, variation_strength, variation_seed, denoise): + if main_seed == variation_seed: + variation_seed += 1 + + end_at_step = steps #min(steps, end_at_step) + start_at_step = round(end_at_step - end_at_step * denoise) + + force_full_denoise = True + disable_noise = True + + device = comfy.model_management.get_torch_device() + + # Generate base noise batch_size, _, height, width = latent_image["samples"].shape + generator = torch.manual_seed(main_seed) base_noise = torch.randn((1, 4, height, width), dtype=torch.float32, device="cpu", generator=generator).repeat(batch_size, 1, 1, 1).cpu() + # Generate variation noise generator = torch.manual_seed(variation_seed) variation_noise = torch.randn((batch_size, 4, height, width), dtype=torch.float32, device="cpu", generator=generator).cpu() slerp_noise = slerp(variation_strength, base_noise, variation_noise) - device = comfy.model_management.get_torch_device() - end_at_step = steps #min(steps, end_at_step) - start_at_step = 0 #min(start_at_step, end_at_step) - real_model = None + # Calculate sigma comfy.model_management.load_model_gpu(model) real_model = model.model sampler = comfy.samplers.KSampler(real_model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler, denoise=1.0, model_options=model.model_options) @@ -1137,11 +1176,11 @@ class KSamplerVariationsWithNoise: work_latent = latent_image.copy() work_latent["samples"] = latent_image["samples"].clone() + slerp_noise * sigma - force_full_denoise = True - #if return_with_leftover_noise == "enable": - # force_full_denoise = False - - disable_noise = True + # if there's a mask we need to expand it to avoid artifacts, 5 pixels should be enough + if "noise_mask" in latent_image: + noise_mask = prepare_mask(latent_image["noise_mask"], latent_image['samples'].shape) + work_latent["samples"] = noise_mask * work_latent["samples"] + (1-noise_mask) * latent_image["samples"] + work_latent['noise_mask'] = expand_mask(latent_image["noise_mask"].clone(), 5, True) return common_ksampler(model, main_seed, steps, cfg, sampler_name, scheduler, positive, negative, work_latent, denoise=1.0, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)