diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py index a3585e716..099334d82 100644 --- a/modules/processing_helpers.py +++ b/modules/processing_helpers.py @@ -29,28 +29,83 @@ def is_refiner_enabled(p): return p.enable_hr and (p.refiner_steps > 0) and (p.refiner_start > 0) and (p.refiner_start < 1) and (shared.sd_refiner is not None) +class ColorCorrectionRef: + __slots__ = ('image', 'lab') + def __init__(self, lab, image): + self.lab = lab + self.image = image + + def setup_color_correction(image): debug("Calibrating color correction") - correction_target = cv2.cvtColor(np.asarray(image.copy()), cv2.COLOR_RGB2LAB) - return correction_target + lab = cv2.cvtColor(np.asarray(image.copy()), cv2.COLOR_RGB2LAB) + return ColorCorrectionRef(lab, image.copy()) -def apply_color_correction(correction, original_image): +def _apply_histogram(correction, original_image): from installer import install install('scikit-image', quiet=True) install('blendmodes', quiet=True) from skimage import exposure from blendmodes.blend import blendLayers, BlendType - log.debug(f"Applying color correction: correction={correction.shape} image={original_image}") + lab = correction.lab if isinstance(correction, ColorCorrectionRef) else correction + log.debug(f"Applying color correction: method=histogram correction={lab.shape} image={original_image}") np_image = np.asarray(original_image) np_recolor = cv2.cvtColor(np_image, cv2.COLOR_RGB2LAB) - np_match = exposure.match_histograms(np_recolor, correction, channel_axis=2) + np_match = exposure.match_histograms(np_recolor, lab, channel_axis=2) np_output = cv2.cvtColor(np_match, cv2.COLOR_LAB2RGB) image = Image.fromarray(np_output.astype("uint8")) image = blendLayers(image, original_image, BlendType.LUMINOSITY) return image +def _apply_wavelet(correction, original_image): + ref_pil = correction.image if isinstance(correction, ColorCorrectionRef) else original_image + ref = torch.from_numpy(np.asarray(ref_pil).astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0) + gen = torch.from_numpy(np.asarray(original_image).astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0) + if ref.shape[2:] != gen.shape[2:]: + ref = torch.nn.functional.interpolate(ref, size=gen.shape[2:], mode='bilinear', align_corners=False) + kernel = torch.tensor([[1, 2, 1], [2, 4, 2], [1, 2, 1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0) / 16.0 + kernel = kernel.expand(3, -1, -1, -1) + log.debug(f"Applying color correction: method=wavelet levels=5 image={original_image}") + gen_highs = [] + current = gen + for _ in range(5): + low = torch.nn.functional.conv2d(current, kernel, padding=1, groups=3) + gen_highs.append(current - low) + current = low + ref_low = ref + for _ in range(5): + ref_low = torch.nn.functional.conv2d(ref_low, kernel, padding=1, groups=3) + result = ref_low + for high in reversed(gen_highs): + result = result + high + result = result.clamp(0, 1).squeeze(0).permute(1, 2, 0).numpy() + return Image.fromarray((result * 255).astype(np.uint8)) + + +def _apply_adain(correction, original_image): + ref_pil = correction.image if isinstance(correction, ColorCorrectionRef) else original_image + ref = torch.from_numpy(np.asarray(ref_pil).astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0) + gen = torch.from_numpy(np.asarray(original_image).astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0) + if ref.shape[2:] != gen.shape[2:]: + ref = torch.nn.functional.interpolate(ref, size=gen.shape[2:], mode='bilinear', align_corners=False) + log.debug(f"Applying color correction: method=adain image={original_image}") + ref_mean = ref.mean(dim=(2, 3), keepdim=True) + ref_std = ref.std(dim=(2, 3), keepdim=True) + 1e-6 + gen_mean = gen.mean(dim=(2, 3), keepdim=True) + gen_std = gen.std(dim=(2, 3), keepdim=True) + 1e-6 + result = (gen - gen_mean) / gen_std * ref_std + ref_mean + result = result.clamp(0, 1).squeeze(0).permute(1, 2, 0).numpy() + return Image.fromarray((result * 255).astype(np.uint8)) + + +def apply_color_correction(correction, original_image, method='histogram'): + methods = {'histogram': _apply_histogram, 'wavelet': _apply_wavelet, 'adain': _apply_adain} + fn = methods.get(method, _apply_histogram) + return fn(correction, original_image) + + def apply_overlay(image: Image, paste_loc, index, overlays): if overlays is None or index >= len(overlays): return image @@ -161,13 +216,18 @@ def slerp_alt(val, lo, hi): # from https://discuss.pytorch.org/t/help-regarding- def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, seed_resize_from_h=0, seed_resize_from_w=0, p=None): - eta_noise_seed_delta = shared.opts.eta_noise_seed_delta or 0 + eta_noise_seed_delta = (getattr(p, 'eta_noise_seed_delta', None) if p is not None else None) + if eta_noise_seed_delta is None: + eta_noise_seed_delta = shared.opts.eta_noise_seed_delta or 0 + enable_batch_seeds = (getattr(p, 'enable_batch_seeds', None) if p is not None else None) + if enable_batch_seeds is None: + enable_batch_seeds = shared.opts.enable_batch_seeds xs = [] # if we have multiple seeds, this means we are working with batch size>1; this then # enables the generation of additional tensors with noise that the sampler will use during its processing. # Using those pre-generated tensors instead of simple torch.randn allows a batch with seeds [100, 101] to # produce the same images as with two batches [100], [101]. - if p is not None and p.sampler is not None and ((len(seeds) > 1 and shared.opts.enable_batch_seeds) or (eta_noise_seed_delta > 0)): + if p is not None and p.sampler is not None and ((len(seeds) > 1 and enable_batch_seeds) or (eta_noise_seed_delta > 0)): sampler_noises = [[] for _ in range(p.sampler.number_of_needed_noises(p))] else: sampler_noises = None @@ -445,14 +505,17 @@ def calculate_refiner_steps(p): def get_generator(p): - if shared.opts.diffusers_generator_device == "Unset": + gen_device_opt = getattr(p, 'diffusers_generator_device', None) if p is not None else None + if gen_device_opt is None: + gen_device_opt = shared.opts.diffusers_generator_device + if gen_device_opt == "Unset": generator_device = None generator = None elif getattr(p, "generator", None) is not None: - generator_device = devices.cpu if shared.opts.diffusers_generator_device == "CPU" else shared.device + generator_device = devices.cpu if gen_device_opt == "CPU" else shared.device generator = p.generator else: - generator_device = devices.cpu if shared.opts.diffusers_generator_device == "CPU" else shared.device + generator_device = devices.cpu if gen_device_opt == "CPU" else shared.device try: p.seeds = [seed if seed != -1 else get_fixed_seed(seed) for seed in p.seeds if seed] devices.randn(p.seeds[0]) @@ -504,7 +567,8 @@ def save_intermediate(p, latents, suffix): info=create_infotext(p, p.all_prompts, p.all_seeds, p.all_subseeds, [], iteration=p.iteration, position_in_batch=i) decoded = processing_vae.vae_decode(latents=latents, model=shared.sd_model, output_type='pil', vae_type=p.vae_type, width=p.width, height=p.height) for j in range(len(decoded)): - images.save_image(decoded[j], path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix=suffix) + _fmt = p.samples_format if p.samples_format is not None else shared.opts.samples_format + images.save_image(decoded[j], path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=_fmt, info=info, p=p, suffix=suffix) def update_sampler(p, sd_model, second_pass=False): @@ -516,7 +580,16 @@ def update_sampler(p, sd_model, second_pass=False): if sampler is None: log.warning(f'Sampler: "{sampler_selection}" not found') sampler = sd_samplers.all_samplers_map.get("UniPC") - sampler = sd_samplers.create_sampler(sampler.name, sd_model) + sched_override_keys = [ + 'schedulers_prediction_type', 'schedulers_beta_schedule', 'schedulers_timesteps', + 'schedulers_sigma', 'schedulers_use_thresholding', 'schedulers_use_loworder', + 'schedulers_solver_order', 'uni_pc_variant', 'schedulers_beta_start', + 'schedulers_beta_end', 'schedulers_shift', 'schedulers_dynamic_shift', + 'schedulers_base_shift', 'schedulers_max_shift', 'schedulers_rescale_betas', + 'schedulers_timestep_spacing', 'schedulers_timesteps_range', + ] + scheduler_overrides = {k: getattr(p, k) for k in sched_override_keys if getattr(p, k, None) is not None} + sampler = sd_samplers.create_sampler(sampler.name, sd_model, scheduler_overrides=scheduler_overrides) if sampler is None or sampler_selection == 'Default': if second_pass: p.hr_sampler = 'Default'