merge: modules/processing_helpers.py

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'