Merge pull request #4690 from vladmandic/fix/color-grading-and-latent-corrections

Fix/color grading and latent corrections

modules/processing_callbacks.py

@@ -40,7 +40,7 @@ def diffusers_callback_legacy(step: int, timestep: int, latents: torch.FloatTens
         latents = torch.from_numpy(latents)
     shared.state.sampling_step = step
     shared.state.current_latent = latents
-    latents = processing_correction.correction_callback(p, timestep, {'latents': latents})
+    latents = processing_correction.correction_callback(p, timestep, {'latents': latents}, step=step)
     if shared.state.interrupted or shared.state.skipped:
         raise AssertionError('Interrupted...')
     if shared.state.paused:
@@ -93,7 +93,7 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict = No
             debug_callback(f"Callback: IP Adapter scales={ip_adapter_scales}")
             pipe.set_ip_adapter_scale(ip_adapter_scales)
     if step != getattr(pipe, 'num_timesteps', 0):
-        kwargs = processing_correction.correction_callback(p, timestep, kwargs, pipe=pipe, initial=step == 0)
+        kwargs = processing_correction.correction_callback(p, timestep, kwargs, pipe=pipe, initial=step == 0, step=step)
     kwargs = prompt_callback(step, kwargs)  # monkey patch for diffusers callback issues
 
     if step == 0:

modules/processing_correction.py

@@ -80,29 +80,64 @@ def color_adjust(tensor, colorstr, ratio):
     return tensor
 
 
-def correction(p, timestep, latent):
-    if timestep > 950 and p.hdr_clamp:
+def correction(p, timestep, latent, step=0):
+    total = getattr(p, 'correction_total_steps', 0)
+    if total > 0:
+        progress = step / total  # 0.0 = first step, ~1.0 = last step
+        is_early = progress < 0.05
+        is_mid = 0.2 <= progress <= 0.7
+        is_late = progress >= 0.8
+        is_sharpen = progress >= 0.7
+        is_very_late = progress >= 0.9
+    else:
+        # fallback to timestep-based ranges for non-flow-match schedulers
+        is_early = timestep > 950
+        is_mid = 600 < timestep < 900
+        is_late = timestep < 200
+        is_sharpen = timestep < 350
+        is_very_late = 1 < timestep < 100
+    if is_early and p.hdr_clamp:
         latent = soft_clamp_tensor(latent, threshold=p.hdr_threshold, boundary=p.hdr_boundary)
         p.extra_generation_params["Latent clamp"] = f'{p.hdr_threshold}/{p.hdr_boundary}'
-    if 600 < timestep < 900 and p.hdr_color != 0:
+    if is_mid and p.hdr_color != 0:
         n = getattr(p, 'correction_steps_mid', 1)
-        latent[1:] = center_tensor(latent[1:], channel_shift=p.hdr_color / n, full_shift=float(p.hdr_mode))
+        num_channels = latent.shape[0]
+        if num_channels <= 4:
+            # SDXL-style: channel 0 is brightness, channels 1+ are color
+            latent[1:] = center_tensor(latent[1:], channel_shift=p.hdr_color / n, full_shift=float(p.hdr_mode))
+        else:
+            # Multi-channel latents (Flux 2, etc.): apply to all channels
+            latent = center_tensor(latent, channel_shift=p.hdr_color / n, full_shift=float(p.hdr_mode))
         p.extra_generation_params["Latent color"] = f'{p.hdr_color}'
-    if 600 < timestep < 900 and p.hdr_tint_ratio != 0:
+    if is_mid and p.hdr_tint_ratio != 0:
         n = getattr(p, 'correction_steps_mid', 1)
-        latent = color_adjust(latent, p.hdr_color_picker, p.hdr_tint_ratio / n)
+        num_channels = latent.shape[0]
+        if num_channels <= 4:
+            # SDXL-style: TAESD color encoding maps to 4-channel latent space
+            latent = color_adjust(latent, p.hdr_color_picker, p.hdr_tint_ratio / n)
+        else:
+            # Multi-channel latents: apply uniform offset to all channels based on tint ratio
+            latent = center_tensor(latent, full_shift=1.0, offset=p.hdr_tint_ratio / n)
         p.extra_generation_params["Latent tint"] = f'{p.hdr_tint_ratio}'
         p.extra_generation_params["Latent tint color"] = p.hdr_color_picker
-    if timestep < 200 and (p.hdr_brightness != 0):
+    if is_late and p.hdr_brightness != 0:
         n = getattr(p, 'correction_steps_late', 1)
-        latent[0:1] = center_tensor(latent[0:1], full_shift=float(p.hdr_mode), offset=p.hdr_brightness / n)
+        num_channels = latent.shape[0]
+        if num_channels <= 4:
+            # SDXL-style: brightness is in channel 0 (luminance)
+            latent[0:1] = center_tensor(latent[0:1], full_shift=float(p.hdr_mode), offset=p.hdr_brightness / n)
+        else:
+            # Multi-channel latents (Flux 2, etc.): scale intensity to avoid color shifts
+            scale = 1.0 + (p.hdr_brightness / n) * 0.25
+            latent = latent * scale
         p.extra_generation_params["Latent brightness"] = f'{p.hdr_brightness}'
-    if timestep < 350 and p.hdr_sharpen != 0:
-        per_step_ratio = 2 ** (timestep / 250) * p.hdr_sharpen / 16
+    if is_sharpen and p.hdr_sharpen != 0:
+        progress_in_range = (step - int(total * 0.7)) / max(int(total * 0.3), 1) if total > 0 else timestep / 350
+        per_step_ratio = 2 ** (progress_in_range * 1.4) * p.hdr_sharpen / 16
         if abs(per_step_ratio) > 0.01:
             latent = sharpen_tensor(latent, ratio=per_step_ratio)
         p.extra_generation_params["Latent sharpen"] = f'{p.hdr_sharpen}'
-    if 1 < timestep < 100 and p.hdr_maximize:
+    if is_very_late and p.hdr_maximize:
         latent = center_tensor(latent, channel_shift=p.hdr_max_center, full_shift=1.0)
         latent = maximize_tensor(latent, boundary=p.hdr_max_boundary)
         p.extra_generation_params["Latent max"] = f'{p.hdr_max_center}/{p.hdr_max_boundary}'
@@ -176,7 +211,7 @@ def _count_steps_below(pipe, threshold):
     return max(count, 1)
 
 
-def correction_callback(p, timestep, kwargs, pipe=None, initial: bool = False):
+def correction_callback(p, timestep, kwargs, pipe=None, initial: bool = False, step: int = 0):
     if initial:
         if not any([p.hdr_clamp, p.hdr_mode, p.hdr_maximize, p.hdr_sharpen, p.hdr_color, p.hdr_brightness, p.hdr_tint_ratio]):
             p.correction_skip = True
@@ -191,12 +226,20 @@ def correction_callback(p, timestep, kwargs, pipe=None, initial: bool = False):
             return kwargs
         p.correction_skip = False
         p.correction_warned = False
-        if pipe is not None:
+        total = getattr(pipe, 'num_timesteps', 0) if pipe is not None else 0
+        if total > 0:
+            p.correction_total_steps = total
+            p.correction_steps_mid = max(int(total * 0.5), 1)  # 20%-70% range
+            p.correction_steps_late = max(int(total * 0.2), 1)  # last 20%
+        elif pipe is not None:
+            p.correction_total_steps = 0
             p.correction_steps_mid = _count_steps_in_range(pipe, 600, 900)
             p.correction_steps_late = _count_steps_below(pipe, 200)
     elif getattr(p, 'correction_skip', False):
         return kwargs
     latents = kwargs["latents"]
+    if debug_enabled:
+        debug(f'Correction callback: step={step} timestep={timestep} latents_shape={latents.shape} total={getattr(p, "correction_total_steps", "unset")} skip={getattr(p, "correction_skip", "unset")}')
     if len(latents.shape) <= 3:  # packed latent
         if pipe is None:
             if not getattr(p, 'correction_warned', False):
@@ -210,11 +253,11 @@ def correction_callback(p, timestep, kwargs, pipe=None, initial: bool = False):
                 p.correction_warned = True
             return kwargs
         for i in range(unpacked.shape[0]):
-            unpacked[i] = correction(p, timestep, unpacked[i])
+            unpacked[i] = correction(p, timestep, unpacked[i], step=step)
         kwargs["latents"] = _repack_latents(unpacked, pack_type, pipe, p)
     elif len(latents.shape) == 4:  # standard batched latent
         for i in range(latents.shape[0]):
-            latents[i] = correction(p, timestep, latents[i])
+            latents[i] = correction(p, timestep, latents[i], step=step)
             if debug_enabled:
                 debug(f"Full Mean: {latents[i].mean().item()}")
                 debug(f"Channel Means: {latents[i].mean(dim=(-1, -2), keepdim=True).flatten().float().cpu().numpy()}")
@@ -224,7 +267,7 @@ def correction_callback(p, timestep, kwargs, pipe=None, initial: bool = False):
     elif len(latents.shape) == 5 and latents.shape[0] == 1:  # probably animatediff
         latents = latents.squeeze(0).permute(1, 0, 2, 3)
         for i in range(latents.shape[0]):
-            latents[i] = correction(p, timestep, latents[i])
+            latents[i] = correction(p, timestep, latents[i], step=step)
         latents = latents.permute(1, 0, 2, 3).unsqueeze(0)
         kwargs["latents"] = latents
     else:

modules/processing_grading.py

@@ -69,6 +69,11 @@ class GradingParams:
     lut_file: str = ""
     lut_strength: float = 1.0
 
+    def __post_init__(self):
+        for f in fields(self):
+            if f.type is float:
+                setattr(self, f.name, float(getattr(self, f.name)))
+
 
 _defaults = GradingParams()
 
@@ -112,18 +117,22 @@ def _apply_shadows_midtones_highlights(img: torch.Tensor, shadows: float, midton
     kornia = _ensure_kornia()
     lab = kornia.color.rgb_to_lab(img)
     L = lab[:, 0:1, :, :] / 100.0  # normalize to [0, 1]
+    strength = 2.0  # scale slider values for more visible effect
     if shadows != 0:
+        s = shadows * strength
         shadow_mask = (1.0 - L).clamp(0, 1) ** 2
-        gamma = 1.0 / (1.0 + shadows) if shadows > 0 else 1.0 - shadows
+        gamma = 1.0 / (1.0 + s) if s > 0 else 1.0 - s
         L = L + shadow_mask * (L.clamp(min=1e-6) ** gamma - L)
     if highlights != 0:
+        h = highlights * strength
         highlight_mask = L.clamp(0, 1) ** 2
-        gamma = 1.0 / (1.0 + highlights) if highlights > 0 else 1.0 - highlights
+        gamma = 1.0 / (1.0 + h) if h > 0 else 1.0 - h
         L = L + highlight_mask * (L.clamp(min=1e-6) ** gamma - L)
     if midtones != 0:
+        m = midtones * strength
         mid_mask = 1.0 - 2.0 * (L - 0.5).abs()
         mid_mask = mid_mask.clamp(0, 1) ** 2
-        gamma = 1.0 / (1.0 + midtones) if midtones > 0 else 1.0 - midtones
+        gamma = 1.0 / (1.0 + m) if m > 0 else 1.0 - m
         L = L + mid_mask * (L.clamp(min=1e-6) ** gamma - L)
     lab[:, 0:1, :, :] = L.clamp(0, 1) * 100.0
     return kornia.color.lab_to_rgb(lab).clamp(0, 1)
@@ -207,7 +216,7 @@ def grade_image(image: Image.Image, params: GradingParams) -> Image.Image:
     if params.gamma != 1.0:
         tensor = kornia.enhance.adjust_gamma(tensor, params.gamma)
     if params.sharpness != 0:
-        tensor = kornia.enhance.sharpness(tensor, params.sharpness)
+        tensor = kornia.enhance.sharpness(tensor, 1.0 + params.sharpness * 4.0)
     if params.color_temp != 6500:
         tensor = _apply_color_temp(tensor, params.color_temp)
 
@@ -215,7 +224,11 @@ def grade_image(image: Image.Image, params: GradingParams) -> Image.Image:
     if params.shadows != 0 or params.midtones != 0 or params.highlights != 0:
         tensor = _apply_shadows_midtones_highlights(tensor, params.shadows, params.midtones, params.highlights)
     if params.clahe_clip > 0:
-        tensor = kornia.enhance.equalize_clahe(tensor, clip_limit=params.clahe_clip, grid_size=(params.clahe_grid, params.clahe_grid))
+        lab = kornia.color.rgb_to_lab(tensor)
+        L = lab[:, 0:1, :, :] / 100.0
+        L = kornia.enhance.equalize_clahe(L, clip_limit=params.clahe_clip, grid_size=(params.clahe_grid, params.clahe_grid))
+        lab[:, 0:1, :, :] = L * 100.0
+        tensor = kornia.color.lab_to_rgb(lab).clamp(0, 1)
 
     # split toning
     if params.shadows_tint != "#000000" or params.highlights_tint != "#ffffff":

modules/ui_sections.py

@@ -168,9 +168,9 @@ def create_latent_inputs(tab):
                 hdr_mode = gr.Dropdown(label="Correction mode", choices=["Relative values", "Absolute values"], type="index", value="Relative values", elem_id=f"{tab}_hdr_mode", show_label=False)
                 hdr_apply_hires = gr.Checkbox(label="Apply to hires", value=True, elem_id=f"{tab}_hdr_apply_hires")
             with gr.Row(elem_id=f"{tab}_correction_row"):
-                hdr_brightness = gr.Slider(minimum=-1.0, maximum=1.0, step=0.05, value=0, label="Latent brightness", elem_id=f"{tab}_hdr_brightness")
-                hdr_sharpen = gr.Slider(minimum=-1.0, maximum=1.0, step=0.05, value=0, label="Latent sharpen", elem_id=f"{tab}_hdr_sharpen")
-                hdr_color = gr.Slider(minimum=0.0, maximum=4.0, step=0.1, value=0.0, label="Latent color", elem_id=f"{tab}_hdr_color")
+                hdr_brightness = gr.Slider(minimum=-4.0, maximum=4.0, step=0.05, value=0, label="Latent brightness", elem_id=f"{tab}_hdr_brightness")
+                hdr_sharpen = gr.Slider(minimum=-4.0, maximum=4.0, step=0.05, value=0, label="Latent sharpen", elem_id=f"{tab}_hdr_sharpen")
+                hdr_color = gr.Slider(minimum=0.0, maximum=16.0, step=0.1, value=0.0, label="Latent color", elem_id=f"{tab}_hdr_color")
             with gr.Row(elem_id=f"{tab}_hdr_clamp_row"):
                 hdr_clamp = gr.Checkbox(label="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")
@@ -181,7 +181,7 @@ def create_latent_inputs(tab):
                 hdr_max_boundary = gr.Slider(minimum=0.5, maximum=2.0, step=0.1, value=1.0, label="Max range", elem_id=f"{tab}_hdr_max_boundary")
             with gr.Row(elem_id=f"{tab}_hdr_color_row"):
                 hdr_color_picker = gr.ColorPicker(label="Tint color", show_label=True, container=False, value=None, elem_id=f"{tab}_hdr_color_picker")
-                hdr_tint_ratio = gr.Slider(label="Tint strength", minimum=-1.0, maximum=1.0, step=0.05, value=0.0, elem_id=f"{tab}_hdr_tint_ratio")
+                hdr_tint_ratio = gr.Slider(label="Tint strength", minimum=-4.0, maximum=4.0, step=0.05, value=0.0, elem_id=f"{tab}_hdr_tint_ratio")
         return hdr_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundary, hdr_color_picker, hdr_tint_ratio, hdr_apply_hires
 
 
@@ -205,7 +205,7 @@ def create_color_inputs(tab):
                 grading_midtones = gr.Slider(minimum=-1.0, maximum=1.0, step=0.05, value=0, label='Midtones', elem_id=f"{tab}_grading_midtones")
                 grading_highlights = gr.Slider(minimum=-1.0, maximum=1.0, step=0.05, value=0, label='Highlights', elem_id=f"{tab}_grading_highlights")
             with gr.Row(elem_id=f"{tab}_grading_clahe_row"):
-                grading_clahe_clip = gr.Slider(minimum=0.0, maximum=40.0, step=1.0, value=0, label='CLAHE clip', elem_id=f"{tab}_grading_clahe_clip")
+                grading_clahe_clip = gr.Slider(minimum=0.0, maximum=5.0, step=0.25, value=0, label='CLAHE clip', elem_id=f"{tab}_grading_clahe_clip")
                 grading_clahe_grid = gr.Slider(minimum=2, maximum=16, step=1, value=8, label='CLAHE grid', elem_id=f"{tab}_grading_clahe_grid")
         with gr.Group():
             with gr.Row(elem_id=f"{tab}_grading_split_row"):