From 5c9ca793e9efcaa77b5287804b0d071507411ffe Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Thu, 1 Feb 2024 18:36:40 -0600
Subject: [PATCH] First Callback Commit

---
 modules/control/run.py           |  6 ++---
 modules/img2img.py               |  7 ++++--
 modules/processing_class.py      | 10 +++++---
 modules/processing_correction.py | 41 ++++++++++++++++++++++++++------
 modules/processing_diffusers.py  | 17 +++++++++----
 modules/shared.py                |  1 +
 modules/txt2img.py               |  8 ++++---
 modules/ui_control.py            |  5 ++--
 modules/ui_img2img.py            |  7 +++---
 modules/ui_sections.py           | 21 ++++++++++++----
 modules/ui_txt2img.py            |  7 +++---
 scripts/xyz_grid.py              |  8 +++++--
 12 files changed, 101 insertions(+), 37 deletions(-)

diff --git a/modules/control/run.py b/modules/control/run.py
index 12a301008..2b3144d91 100644
--- a/modules/control/run.py
+++ b/modules/control/run.py
@@ -35,7 +35,7 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
                 prompt, negative, styles, steps, sampler_index,
                 seed, subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w,
                 cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, full_quality, restore_faces, tiling,
-                hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry,
+                hdr_clamp, hdr_boundary, hdr_threshold, hdr_brightness, hdr_center, hdr_color_correction, hdr_maximize, hdr_max_center, hdr_max_boundry,
                 resize_mode_before, resize_name_before, width_before, height_before, scale_by_before, selected_scale_tab_before,
                 resize_mode_after, resize_name_after, width_after, height_after, scale_by_after, selected_scale_tab_after,
                 denoising_strength, batch_count, batch_size,
@@ -81,8 +81,8 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
         hdr_boundary = hdr_boundary,
         hdr_threshold = hdr_threshold,
         hdr_center = hdr_center,
-        hdr_channel_shift = hdr_channel_shift,
-        hdr_full_shift = hdr_full_shift,
+        hdr_color_correction = hdr_color_correction,
+        hdr_brightness = hdr_brightness,
         hdr_maximize = hdr_maximize,
         hdr_max_center = hdr_max_center,
         hdr_max_boundry = hdr_max_boundry,
diff --git a/modules/img2img.py b/modules/img2img.py
index d1c2ae0bd..38a7cf5c3 100644
--- a/modules/img2img.py
+++ b/modules/img2img.py
@@ -124,6 +124,7 @@ def img2img(id_task: str, mode: int,
             cfg_scale, image_cfg_scale,
             diffusers_guidance_rescale,
             sag_scale,
+            cfg_end,
             refiner_start,
             clip_skip,
             denoising_strength,
@@ -134,7 +135,7 @@ def img2img(id_task: str, mode: int,
             resize_mode, resize_name,
             inpaint_full_res, inpaint_full_res_padding, inpainting_mask_invert,
             img2img_batch_files, img2img_batch_input_dir, img2img_batch_output_dir, img2img_batch_inpaint_mask_dir,
-            hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry,
+            hdr_clamp, hdr_boundary, hdr_threshold, hdr_brightness, hdr_center, hdr_color_correction, hdr_sharpen, hdr_sharpen_ratio, hdr_sharpen_start, hdr_maximize, hdr_max_center, hdr_max_boundry,
             override_settings_texts,
             *args): # pylint: disable=unused-argument
 
@@ -217,6 +218,7 @@ def img2img(id_task: str, mode: int,
         n_iter=n_iter,
         steps=steps,
         cfg_scale=cfg_scale,
+        cfg_end=cfg_end,
         clip_skip=clip_skip,
         width=width,
         height=height,
@@ -238,7 +240,8 @@ def img2img(id_task: str, mode: int,
         inpaint_full_res_padding=inpaint_full_res_padding,
         inpainting_mask_invert=inpainting_mask_invert,
         hdr_clamp=hdr_clamp, hdr_boundary=hdr_boundary, hdr_threshold=hdr_threshold,
-        hdr_center=hdr_center, hdr_channel_shift=hdr_channel_shift, hdr_full_shift=hdr_full_shift,
+        hdr_center=hdr_center, hdr_color_correction=hdr_color_correction, hdr_brightness=hdr_brightness,
+        hdr_sharpen=hdr_sharpen, hdr_sharpen_ratio=hdr_sharpen_ratio, hdr_sharpen_start=hdr_sharpen_start,
         hdr_maximize=hdr_maximize, hdr_max_center=hdr_max_center, hdr_max_boundry=hdr_max_boundry,
         override_settings=override_settings,
     )
diff --git a/modules/processing_class.py b/modules/processing_class.py
index db712dcaa..4cfb27eb7 100644
--- a/modules/processing_class.py
+++ b/modules/processing_class.py
@@ -16,7 +16,7 @@ class StableDiffusionProcessing:
     """
     The first set of paramaters: sd_models -> do_not_reload_embeddings represent the minimum required to create a StableDiffusionProcessing
     """
-    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt: str = "", styles: List[str] = None, seed: int = -1, subseed: int = -1, subseed_strength: float = 0, seed_resize_from_h: int = -1, seed_resize_from_w: int = -1, seed_enable_extras: bool = True, sampler_name: str = None, hr_sampler_name: str = None, batch_size: int = 1, n_iter: int = 1, steps: int = 50, cfg_scale: float = 7.0, image_cfg_scale: float = None, clip_skip: int = 1, width: int = 512, height: int = 512, full_quality: bool = True, restore_faces: bool = False, tiling: bool = False, do_not_save_samples: bool = False, do_not_save_grid: bool = False, extra_generation_params: Dict[Any, Any] = None, overlay_images: Any = None, negative_prompt: str = None, eta: float = None, do_not_reload_embeddings: bool = False, denoising_strength: float = 0, diffusers_guidance_rescale: float = 0.7, sag_scale: float = 0.0, resize_mode: int = 0, resize_name: str = 'None', scale_by: float = 0, selected_scale_tab: int = 0, hdr_clamp: bool = False, hdr_boundary: float = 4.0, hdr_threshold: float = 3.5, hdr_center: bool = False, hdr_channel_shift: float = 0.8, hdr_full_shift: float = 0.8, hdr_maximize: bool = False, hdr_max_center: float = 0.6, hdr_max_boundry: float = 1.0, override_settings: Dict[str, Any] = None, override_settings_restore_afterwards: bool = True, sampler_index: int = None, script_args: list = None): # pylint: disable=unused-argument
+    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt: str = "", styles: List[str] = None, seed: int = -1, subseed: int = -1, subseed_strength: float = 0, seed_resize_from_h: int = -1, seed_resize_from_w: int = -1, seed_enable_extras: bool = True, sampler_name: str = None, hr_sampler_name: str = None, batch_size: int = 1, n_iter: int = 1, steps: int = 50, cfg_scale: float = 7.0, image_cfg_scale: float = None, clip_skip: int = 1, width: int = 512, height: int = 512, full_quality: bool = True, restore_faces: bool = False, tiling: bool = False, do_not_save_samples: bool = False, do_not_save_grid: bool = False, extra_generation_params: Dict[Any, Any] = None, overlay_images: Any = None, negative_prompt: str = None, eta: float = None, do_not_reload_embeddings: bool = False, denoising_strength: float = 0, diffusers_guidance_rescale: float = 0.7, sag_scale: float = 0.0, cfg_end: float = 1, resize_mode: int = 0, resize_name: str = 'None', scale_by: float = 0, selected_scale_tab: int = 0, hdr_clamp: bool = False, hdr_boundary: float = 4.0, hdr_threshold: float = 3.5, hdr_center: bool = False, hdr_color_correction: float = 0.8, hdr_brightness: float = 0.8, hdr_sharpen: bool = False, hdr_sharpen_ratio: float = 0.0, hdr_sharpen_start: int = 500, hdr_maximize: bool = False, hdr_max_center: float = 0.6, hdr_max_boundry: float = 1.0, override_settings: Dict[str, Any] = None, override_settings_restore_afterwards: bool = True, sampler_index: int = None, script_args: list = None): # pylint: disable=unused-argument
         self.outpath_samples: str = outpath_samples
         self.outpath_grids: str = outpath_grids
         self.prompt: str = prompt
@@ -39,6 +39,7 @@ class StableDiffusionProcessing:
         self.image_cfg_scale = image_cfg_scale
         self.diffusers_guidance_rescale = diffusers_guidance_rescale
         self.sag_scale = sag_scale
+        self.cfg_end = cfg_end
         if devices.backend == "ipex" and width == 1024 and height == 1024 and not torch.xpu.has_fp64_dtype() and os.environ.get('DISABLE_IPEX_1024_WA', None) is None:
             width = 1080
             height = 1080
@@ -118,8 +119,11 @@ class StableDiffusionProcessing:
         self.hdr_boundary = hdr_boundary
         self.hdr_threshold = hdr_threshold
         self.hdr_center = hdr_center
-        self.hdr_channel_shift = hdr_channel_shift
-        self.hdr_full_shift = hdr_full_shift
+        self.hdr_color_correction = hdr_color_correction
+        self.hdr_brightness = hdr_brightness
+        self.hdr_sharpen = hdr_sharpen
+        self.hdr_sharpen_ratio = hdr_sharpen_ratio
+        self.hdr_sharpen_start = hdr_sharpen_start
         self.hdr_maximize = hdr_maximize
         self.hdr_max_center = hdr_max_center
         self.hdr_max_boundry = hdr_max_boundry
diff --git a/modules/processing_correction.py b/modules/processing_correction.py
index 2f842a596..e077ea173 100644
--- a/modules/processing_correction.py
+++ b/modules/processing_correction.py
@@ -12,6 +12,21 @@ debug = shared.log.trace if os.environ.get('SD_HDR_DEBUG', None) is not None els
 debug('Trace: HDR')
 
 
+def sharpen_tensor(tensor, ratio=0):
+    if ratio == 0:
+        print("early exit...")
+        return tensor
+    kernel = torch.ones((3, 3), dtype=tensor.dtype, device=tensor.device)
+    kernel[1, 1] = 5.0
+    kernel /= kernel.sum()
+    kernel = kernel.expand(tensor.shape[-3], 1, kernel.shape[0], kernel.shape[1])
+    result_tmp = torch.nn.functional.conv2d(tensor, kernel, groups=tensor.shape[-3])
+    result = tensor.clone()
+    result[..., 1:-1, 1:-1] = result_tmp
+    output = (1.0 + ratio) * tensor + (0 - ratio) * result
+    return soft_clamp_tensor(output, threshold=0.95)
+
+
 def soft_clamp_tensor(tensor, threshold=0.8, boundary=4):
     # shrinking towards the mean; will also remove outliers
     if max(abs(tensor.max()), abs(tensor.min())) < boundary or threshold == 0:
@@ -29,11 +44,12 @@ def soft_clamp_tensor(tensor, threshold=0.8, boundary=4):
     return tensor
 
 
-def center_tensor(tensor, channel_shift=1.0, full_shift=1.0):
-    if channel_shift == 0 and full_shift == 0:
+def center_tensor(tensor, channel_shift=0.0, full_shift=0.0, offset=0.0):
+    if channel_shift == 0 and full_shift == 0 and offset == 0:
         return tensor
+    debug(f'HDR center: Before Adjustment: Full mean={tensor.mean().item()} Channel means={tensor.mean(dim=(-1, -2)).cpu().numpy()}')
     tensor -= tensor.mean(dim=(-1, -2), keepdim=True) * channel_shift
-    tensor -= tensor.mean() * full_shift
+    tensor -= tensor.mean() * full_shift - offset
     debug(f'HDR center: channel-shift={channel_shift} full-shift={full_shift}')
     debug(f'HDR center: After Adjustment: Full mean={tensor.mean().item()} Channel means={tensor.mean(dim=(-1, -2)).cpu().numpy()}')
     return tensor
@@ -55,9 +71,20 @@ def correction(p, timestep, latent):
     if timestep > 950 and p.hdr_clamp:
         p.extra_generation_params["HDR clamp"] = f'{p.hdr_threshold}/{p.hdr_boundary}'
         latent = soft_clamp_tensor(latent, threshold=p.hdr_threshold, boundary=p.hdr_boundary)
-    if timestep > 700 and p.hdr_center:
-        p.extra_generation_params["HDR center"] = f'{p.hdr_channel_shift}/{p.hdr_full_shift}'
-        latent = center_tensor(latent, channel_shift=p.hdr_channel_shift, full_shift=p.hdr_full_shift)
+    if 500 < timestep < 800 and (p.hdr_center or p.hdr_brightness):
+        p.extra_generation_params["HDR center"] = f'{p.hdr_color_correction}/{p.hdr_brightness}'
+        latent[0:1] = center_tensor(latent[0:1], full_shift=float(p.hdr_center), offset=p.hdr_brightness)  # Brightness
+        p.hdr_center = 0
+        p.hdr_brightness = 0
+    if 500 < timestep < 800 and p.hdr_color_correction != 0:
+        latent[1:] = center_tensor(latent[1:], channel_shift=p.hdr_color_correction, full_shift=0.0)  # Color
+        p.hdr_color_correction = 0
+    if timestep < p.hdr_sharpen_start and p.hdr_sharpen:
+        print(f"Sharpening... {p.hdr_sharpen_ratio}/{p.hdr_sharpen_start}")
+        p.extra_generation_params["HDR sharpen"] = f'{p.hdr_sharpen_ratio}/{p.hdr_sharpen_start}'
+        latent = sharpen_tensor(latent, ratio=p.hdr_sharpen_ratio)
+        # p.hdr_sharpen = False
+        p.hdr_sharpen_ratio *= 0.5 * 2**0.5
     if 1 < timestep < 100 and p.hdr_maximize:
         p.extra_generation_params["HDR max"] = f'{p.hdr_max_center}/{p.hdr_max_boundry}'
         latent = center_tensor(latent, channel_shift=p.hdr_max_center, full_shift=1.0)
@@ -66,7 +93,7 @@ def correction(p, timestep, latent):
 
 
 def correction_callback(p, timestep, kwargs):
-    if not p.hdr_clamp and not p.hdr_center and not p.hdr_maximize:
+    if not any([p.hdr_clamp, p.hdr_center, p.hdr_maximize, p.hdr_sharpen, p.hdr_color_correction, p.hdr_brightness]):
         return kwargs
     latents = kwargs["latents"]
     debug('')
diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py
index 482da9a71..a137ed46d 100644
--- a/modules/processing_diffusers.py
+++ b/modules/processing_diffusers.py
@@ -50,8 +50,12 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
                     raise AssertionError('Interrupted...')
                 time.sleep(0.1)
 
-    def diffusers_callback(_pipe, step: int, timestep: int, kwargs: dict):
+    def diffusers_callback(pipe, step: int, timestep: int, kwargs: dict):
         shared.state.sampling_step = step
+        if shared.opts.nan_skip:
+            latents = kwargs.get('latents', None)
+            if latents is not None:
+                assert not torch.isnan(latents[..., 0, 0]).all(), f'NaN detected at step {step}: Skipping...'
         if shared.state.interrupted or shared.state.skipped:
             raise AssertionError('Interrupted...')
         if shared.state.paused:
@@ -62,7 +66,8 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
                 time.sleep(0.1)
         if kwargs.get('latents', None) is None:
             return kwargs
-        kwargs = processing_correction.correction_callback(p, timestep, kwargs)
+        if step != pipe.num_timesteps:
+            kwargs = processing_correction.correction_callback(p, timestep, kwargs)
         if p.scheduled_prompt and 'prompt_embeds' in kwargs and 'negative_prompt_embeds' in kwargs:
             try:
                 i = (step + 1) % len(p.prompt_embeds)
@@ -71,6 +76,10 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
                 kwargs["negative_prompt_embeds"] = p.negative_embeds[j][0:1].expand(kwargs["negative_prompt_embeds"].shape)
             except Exception as e:
                 shared.log.debug(f"Callback: {e}")
+        if step == int(pipe.num_timesteps * p.cfg_end) and 'prompt_embeds' in kwargs and 'negative_prompt_embeds' in kwargs:
+            pipe._guidance_scale = 0.0
+            for key in {"prompt_embeds", "negative_prompt_embeds", "add_text_embeds", "add_time_ids"} & set(kwargs):
+                kwargs[key] = kwargs[key].chunk(2)[-1]
         shared.state.current_latent = kwargs['latents']
         if shared.cmd_opts.profile and shared.profiler is not None:
             shared.profiler.step()
@@ -202,7 +211,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
         elif 'callback_on_step_end_tensor_inputs' in possible:
             args['callback_on_step_end'] = diffusers_callback
             if 'prompt_embeds' in possible and 'negative_prompt_embeds' in possible:
-                args['callback_on_step_end_tensor_inputs'] = ['latents', 'prompt_embeds', 'negative_prompt_embeds']
+                args['callback_on_step_end_tensor_inputs'] = model._callback_tensor_inputs
             else:
                 args['callback_on_step_end_tensor_inputs'] = ['latents']
         for arg in kwargs:
@@ -259,7 +268,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
         if p.hdr_clamp or p.hdr_center or p.hdr_maximize:
             txt = 'HDR:'
             txt += f' Clamp threshold={p.hdr_threshold} boundary={p.hdr_boundary}' if p.hdr_clamp else ' Clamp off'
-            txt += f' Center channel-shift={p.hdr_channel_shift} full-shift={p.hdr_full_shift}' if p.hdr_center else ' Center off'
+            txt += f' Center channel-shift={p.hdr_color_correction} full-shift={p.hdr_brightness}' if p.hdr_center else ' Center off'
             txt += f' Maximize boundary={p.hdr_max_boundry} center={p.hdr_max_center}' if p.hdr_maximize else ' Maximize off'
             shared.log.debug(txt)
         # components = [{ k: getattr(v, 'device', None) } for k, v in model.components.items()]
diff --git a/modules/shared.py b/modules/shared.py
index b7389b499..d41acd72c 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -345,6 +345,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
     "upcast_attn": OptionInfo(False, "Upcast attention layer"),
     "cuda_cast_unet": OptionInfo(False, "Fixed UNet precision"),
     "disable_nan_check": OptionInfo(True, "Disable NaN check", gr.Checkbox, {"visible": False}),
+    "nan_skip": OptionInfo(False, "Skip Generation if NaN found in latents", gr.Checkbox, {"visible": True}),
     "rollback_vae": OptionInfo(False, "Attempt VAE roll back for NaN values"),
 
     "cross_attention_sep": OptionInfo("<h2>Attention</h2>", "", gr.HTML),
diff --git a/modules/txt2img.py b/modules/txt2img.py
index cbd14c00d..f2a9f3416 100644
--- a/modules/txt2img.py
+++ b/modules/txt2img.py
@@ -13,14 +13,14 @@ def txt2img(id_task,
             steps, sampler_index, hr_sampler_index,
             full_quality, restore_faces, tiling,
             n_iter, batch_size,
-            cfg_scale, image_cfg_scale, diffusers_guidance_rescale, sag_scale,
+            cfg_scale, image_cfg_scale, diffusers_guidance_rescale, sag_scale, cfg_end,
             clip_skip,
             seed, subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w,
             height, width,
             enable_hr, denoising_strength,
             hr_scale, hr_upscaler, hr_force, hr_second_pass_steps, hr_resize_x, hr_resize_y,
             refiner_steps, refiner_start, refiner_prompt, refiner_negative,
-            hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry,
+            hdr_clamp, hdr_boundary, hdr_threshold, hdr_brightness, hdr_center, hdr_color_correction, hdr_sharpen, hdr_sharpen_ratio, hdr_sharpen_start, hdr_maximize, hdr_max_center, hdr_max_boundry,
             override_settings_texts,
             *args):
 
@@ -58,6 +58,7 @@ def txt2img(id_task,
         image_cfg_scale=image_cfg_scale,
         diffusers_guidance_rescale=diffusers_guidance_rescale,
         sag_scale=sag_scale,
+        cfg_end=cfg_end,
         clip_skip=clip_skip,
         width=width,
         height=height,
@@ -77,7 +78,8 @@ def txt2img(id_task,
         refiner_prompt=refiner_prompt,
         refiner_negative=refiner_negative,
         hdr_clamp=hdr_clamp, hdr_boundary=hdr_boundary, hdr_threshold=hdr_threshold,
-        hdr_center=hdr_center, hdr_channel_shift=hdr_channel_shift, hdr_full_shift=hdr_full_shift,
+        hdr_center=hdr_center, hdr_color_correction=hdr_color_correction, hdr_brightness=hdr_brightness,
+        hdr_sharpen=hdr_sharpen, hdr_sharpen_ratio=hdr_sharpen_ratio, hdr_sharpen_start=hdr_sharpen_start,
         hdr_maximize=hdr_maximize, hdr_max_center=hdr_max_center, hdr_max_boundry=hdr_max_boundry,
         override_settings=override_settings,
     )
diff --git a/modules/ui_control.py b/modules/ui_control.py
index 0b58ba857..004df5ac9 100644
--- a/modules/ui_control.py
+++ b/modules/ui_control.py
@@ -103,7 +103,8 @@ def create_ui(_blocks: gr.Blocks=None):
 
                 mask_controls = masking.create_segment_ui()
 
-                cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry = ui_sections.create_advanced_inputs('control')
+                cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, cfg_end, full_quality, restore_faces, tiling= ui_sections.create_advanced_inputs('control')
+                hdr_clamp, hdr_boundary, hdr_threshold, hdr_brightness, hdr_center, hdr_color_correction, hdr_sharpen, hdr_sharpen_ratio, hdr_sharpen_start, hdr_maximize, hdr_max_center, hdr_max_boundry = ui_sections.create_callback_inputs('control')
 
                 with gr.Accordion(open=False, label="Video", elem_id="control_video", elem_classes=["small-accordion"]):
                     with gr.Row():
@@ -487,7 +488,7 @@ def create_ui(_blocks: gr.Blocks=None):
                     prompt, negative, styles,
                     steps, sampler_index,
                     seed, subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w,
-                    cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry,
+                    cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, cfg_end, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_brightness, hdr_center, hdr_color_correction, hdr_sharpen, hdr_sharpen_ratio, hdr_sharpen_start, hdr_maximize, hdr_max_center, hdr_max_boundry,
                     resize_mode_before, resize_name_before, width_before, height_before, scale_by_before, selected_scale_tab_before,
                     resize_mode_after, resize_name_after, width_after, height_after, scale_by_after, selected_scale_tab_after,
                     denoising_strength, batch_count, batch_size,
diff --git a/modules/ui_img2img.py b/modules/ui_img2img.py
index eabbf778f..957aee48b 100644
--- a/modules/ui_img2img.py
+++ b/modules/ui_img2img.py
@@ -130,7 +130,8 @@ def create_ui():
                             denoising_strength = gr.Slider(minimum=0.0, maximum=0.99, step=0.01, label='Denoising strength', value=0.50, elem_id="img2img_denoising_strength")
                             refiner_start = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Denoise start', value=0.0, elem_id="img2img_refiner_start")
 
-                    cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry = ui_sections.create_advanced_inputs('img2img')
+                    cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, cfg_end, full_quality, restore_faces, tiling = ui_sections.create_advanced_inputs('img2img')
+                    hdr_clamp, hdr_boundary, hdr_threshold, hdr_brightness, hdr_center, hdr_color_correction, hdr_sharpen, hdr_sharpen_ratio, hdr_sharpen_start, hdr_maximize, hdr_max_center, hdr_max_boundry = ui_sections.create_callback_inputs('img2img')
 
                     # with gr.Group(elem_id="inpaint_controls", visible=False) as inpaint_controls:
                     with gr.Accordion(open=True, label="Mask", elem_classes=["small-accordion"], elem_id="img2img_mask_group") as inpaint_controls:
@@ -179,7 +180,7 @@ def create_ui():
                 full_quality, restore_faces, tiling,
                 batch_count, batch_size,
                 cfg_scale, image_cfg_scale,
-                diffusers_guidance_rescale, sag_scale,
+                diffusers_guidance_rescale, sag_scale, cfg_end,
                 refiner_start,
                 clip_skip,
                 denoising_strength,
@@ -190,7 +191,7 @@ def create_ui():
                 resize_mode, resize_name,
                 inpaint_full_res, inpaint_full_res_padding, inpainting_mask_invert,
                 img2img_batch_files, img2img_batch_input_dir, img2img_batch_output_dir, img2img_batch_inpaint_mask_dir,
-                hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry,
+                hdr_clamp, hdr_boundary, hdr_threshold, hdr_brightness, hdr_center, hdr_color_correction, hdr_sharpen, hdr_sharpen_ratio, hdr_sharpen_start, hdr_maximize, hdr_max_center, hdr_max_boundry,
                 override_settings,
             ]
             img2img_dict = dict(
diff --git a/modules/ui_sections.py b/modules/ui_sections.py
index f6456df14..12b240c15 100644
--- a/modules/ui_sections.py
+++ b/modules/ui_sections.py
@@ -107,6 +107,8 @@ def create_advanced_inputs(tab):
             with gr.Row():
                 diffusers_guidance_rescale = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Guidance rescale', value=0.7, elem_id=f"{tab}_image_cfg_rescale", visible=shared.backend == shared.Backend.DIFFUSERS)
                 diffusers_sag_scale = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Self-attention guidance', value=0.0, elem_id=f"{tab}_image_sag_scale", visible=shared.backend == shared.Backend.DIFFUSERS)
+            with gr.Row():
+                cfg_end = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label='CFG end %', value=1.0, elem_id=f"{tab}_cfg_end")
             with gr.Row():
                 clip_skip = gr.Slider(label='CLIP skip', value=1, minimum=1, maximum=14, step=1, elem_id=f"{tab}_clip_skip", interactive=True)
         with gr.Group():
@@ -114,21 +116,30 @@ def create_advanced_inputs(tab):
                 full_quality = gr.Checkbox(label='Full quality', value=True, elem_id=f"{tab}_full_quality")
                 restore_faces = gr.Checkbox(label='Face restore', value=False, visible=len(shared.face_restorers) > 1, elem_id=f"{tab}_restore_faces")
                 tiling = gr.Checkbox(label='Tiling', value=False, elem_id=f"{tab}_tiling", visible=shared.backend == shared.Backend.ORIGINAL)
+
+    return cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, diffusers_sag_scale, cfg_end, full_quality, restore_faces, tiling
+
+def create_callback_inputs(tab):
+    with gr.Accordion(open=False, label="Callback", elem_id=f"{tab}_callback", elem_classes=["small-accordion"], visible=shared.backend == shared.Backend.DIFFUSERS):
         with gr.Group(visible=shared.backend == shared.Backend.DIFFUSERS):
             with gr.Row():
                 hdr_clamp = gr.Checkbox(label='HDR clamp', value=False, elem_id=f"{tab}_hdr_clamp")
                 hdr_boundary = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=4.0,  label='Range', elem_id=f"{tab}_hdr_boundary")
                 hdr_threshold = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.95,  label='Threshold', elem_id=f"{tab}_hdr_threshold")
             with gr.Row():
-                hdr_center = gr.Checkbox(label='HDR center', value=False, elem_id=f"{tab}_hdr_center")
-                hdr_channel_shift = gr.Slider(minimum=0.0, maximum=2.0, step=0.1, value=1.0,  label='Channel shift', elem_id=f"{tab}_hdr_channel_shift")
-                hdr_full_shift = gr.Slider(minimum=0.0, maximum=2.0, step=0.1, value=1,  label='Full shift', elem_id=f"{tab}_hdr_full_shift")
+                hdr_brightness = gr.Slider(minimum=-1.0, maximum=1.0, step=0.1, value=0,  label='Brightness', elem_id=f"{tab}_hdr_brightness")
+                hdr_center = gr.Dropdown(label="HDR Brightness Center", choices=["Relative", "Absolute"], type="index", value="Relative", elem_id=f"{tab}_hdr_center", show_label=False)
+            with gr.Row():
+                hdr_color_correction = gr.Slider(minimum=0.0, maximum=2.0, step=0.1, value=0.0,  label='Color Correction', elem_id=f"{tab}_hdr_color_correction")
+            with gr.Row():
+                hdr_sharpen = gr.Checkbox(label='HDR sharpen', value=False, elem_id=f"{tab}_sharpen")
+                hdr_sharpen_ratio = gr.Slider(minimum=-1.0, maximum=1.0, step=0.1, value=0,  label='Sharpen ratio', elem_id=f"{tab}_hdr_sharpen_ratio")
+                hdr_sharpen_start = gr.Slider(minimum=0, maximum=1000, step=50, value=500,  label='Sharpen start', elem_id=f"{tab}_hdr_sharpen_start")
             with gr.Row():
                 hdr_maximize = gr.Checkbox(label='HDR maximize', value=False, elem_id=f"{tab}_hdr_maximize")
                 hdr_max_center = gr.Slider(minimum=0.0, maximum=2.0, step=0.1, value=0.6,  label='Center', elem_id=f"{tab}_hdr_max_center")
                 hdr_max_boundry = gr.Slider(minimum=0.5, maximum=2.0, step=0.1, value=1.0,  label='Max Range', elem_id=f"{tab}_hdr_max_boundry")
-    return cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, diffusers_sag_scale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry
-
+    return hdr_clamp, hdr_boundary, hdr_threshold, hdr_brightness, hdr_center, hdr_color_correction, hdr_sharpen, hdr_sharpen_ratio, hdr_sharpen_start, hdr_maximize, hdr_max_center, hdr_max_boundry
 
 def create_sampler_and_steps_selection(choices, tabname):
     def set_sampler_original_options(sampler_options, sampler_algo):
diff --git a/modules/ui_txt2img.py b/modules/ui_txt2img.py
index 38177a451..67abf2faf 100644
--- a/modules/ui_txt2img.py
+++ b/modules/ui_txt2img.py
@@ -45,7 +45,8 @@ def create_ui():
                     steps, sampler_index = ui_sections.create_sampler_inputs('txt2img')
                     batch_count, batch_size = ui_sections.create_batch_inputs('txt2img')
                     seed, reuse_seed, subseed, reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w = ui_sections.create_seed_inputs('txt2img')
-                    cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry = ui_sections.create_advanced_inputs('txt2img')
+                    cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, cfg_end, full_quality, restore_faces, tiling = ui_sections.create_advanced_inputs('txt2img')
+                    hdr_clamp, hdr_boundary, hdr_threshold, hdr_brightness, hdr_center, hdr_color_correction, hdr_sharpen, hdr_sharpen_ratio, hdr_sharpen_start, hdr_maximize, hdr_max_center, hdr_max_boundry = ui_sections.create_callback_inputs('txt2img')
                     enable_hr, hr_sampler_index, denoising_strength, hr_final_resolution, hr_upscaler, hr_force, hr_second_pass_steps, hr_scale, hr_resize_x, hr_resize_y, refiner_steps, refiner_start, refiner_prompt, refiner_negative = ui_sections.create_hires_inputs('txt2img')
                     override_settings = ui_common.create_override_inputs('txt2img')
 
@@ -72,14 +73,14 @@ def create_ui():
                 steps, sampler_index, hr_sampler_index,
                 full_quality, restore_faces, tiling,
                 batch_count, batch_size,
-                cfg_scale, image_cfg_scale, diffusers_guidance_rescale, sag_scale,
+                cfg_scale, image_cfg_scale, diffusers_guidance_rescale, sag_scale, cfg_end,
                 clip_skip,
                 seed, subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w,
                 height, width,
                 enable_hr, denoising_strength,
                 hr_scale, hr_upscaler, hr_force, hr_second_pass_steps, hr_resize_x, hr_resize_y,
                 refiner_steps, refiner_start, refiner_prompt, refiner_negative,
-                hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry,
+                hdr_clamp, hdr_boundary, hdr_threshold, hdr_brightness, hdr_center, hdr_color_correction, hdr_sharpen, hdr_sharpen_ratio, hdr_sharpen_start, hdr_maximize, hdr_max_center, hdr_max_boundry,
                 override_settings,
             ]
             txt2img_dict = dict(
diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py
index 39ed66855..9ccef2142 100644
--- a/scripts/xyz_grid.py
+++ b/scripts/xyz_grid.py
@@ -230,6 +230,7 @@ axis_options = [
     AxisOption("Seed", int, apply_field("seed")),
     AxisOption("Steps", int, apply_field("steps")),
     AxisOption("CFG Scale", float, apply_field("cfg_scale")),
+    AxisOption("CFG End", float, apply_field("cfg_end")),
     AxisOption("Variation seed", int, apply_field("subseed")),
     AxisOption("Variation strength", float, apply_field("subseed_strength")),
     AxisOption("Clip skip", int, apply_clip_skip),
@@ -258,9 +259,12 @@ axis_options = [
     AxisOption("[Refiner] Refiner steps", float, apply_field("refiner_steps")),
     AxisOption("[HDR] Clamp boundary", float, apply_field("hdr_boundary")),
     AxisOption("[HDR] Clamp threshold", float, apply_field("hdr_threshold")),
-    AxisOption("[HDR] Center channel shift", float, apply_field("hdr_channel_shift")),
-    AxisOption("[HDR] Center full shift", float, apply_field("hdr_full_shift")),
+    AxisOption("[HDR] Color Correction", float, apply_field("hdr_color_correction")),
+    AxisOption("[HDR] Brightness", float, apply_field("hdr_brightness")),
+    AxisOption("[HDR] Brightness Center", float, apply_field("hdr_center")),
     AxisOption("[HDR] Maximize center shift", float, apply_field("hdr_max_center")),
+    AxisOption("[HDR] Sharpen ratio", float, apply_field("hdr_sharpen_ratio")),
+    AxisOption("[HDR] Sharpen start", float, apply_field("hdr_sharpen_start")),
     AxisOption("[HDR] Maximize boundary", float, apply_field("hdr_max_boundry")),
     AxisOption("[ToMe] Token merging ratio (txt2img)", float, apply_override('token_merging_ratio')),
     AxisOption("[ToMe] Token merging ratio (hires)", float, apply_override('token_merging_ratio_hr')),