From ee2fc75619435da4a047a041199eff9b030e078b Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Thu, 19 Mar 2026 04:00:12 +0000 Subject: [PATCH 1/3] fix(grading): color grading type coercion, CLAHE, sharpness and tone scaling - add __post_init__ to GradingParams to coerce Gradio int returns to float - apply CLAHE to L channel in Lab space instead of per-RGB channel - fix sharpness scaling (kornia factor 1.0=no change, map 0-based slider) - increase shadows/midtones/highlights gamma strength 2x for visible effect - reduce CLAHE clip slider range to 0-5 with finer 0.25 step --- modules/processing_grading.py | 23 ++++++++++++++++++----- modules/ui_sections.py | 2 +- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/modules/processing_grading.py b/modules/processing_grading.py index 903385f1d..003f93f07 100644 --- a/modules/processing_grading.py +++ b/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": diff --git a/modules/ui_sections.py b/modules/ui_sections.py index 326116c12..654c7ee4b 100644 --- a/modules/ui_sections.py +++ b/modules/ui_sections.py @@ -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"): From 423c3ef8b79c2ff073451704d2434296c4292c36 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Thu, 19 Mar 2026 04:03:13 +0000 Subject: [PATCH 2/3] fix(correction): step-based progress and multi-channel latent support replace hardcoded timestep thresholds with step-based progress percentages so corrections work with flow-match schedulers (Flux 2, etc.) adapt brightness, color and tint corrections for multi-channel latents: - brightness uses multiplicative scaling instead of additive offset - color applies to all channels instead of skipping channel 0 - tint falls back to uniform offset when TAESD encoding is unavailable pass step parameter through correction_callback for progress calculation --- modules/processing_callbacks.py | 4 +- modules/processing_correction.py | 75 +++++++++++++++++++++++++------- 2 files changed, 61 insertions(+), 18 deletions(-) diff --git a/modules/processing_callbacks.py b/modules/processing_callbacks.py index bfd036f92..0f6fdb533 100644 --- a/modules/processing_callbacks.py +++ b/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: diff --git a/modules/processing_correction.py b/modules/processing_correction.py index 6fc206063..1c3751bb9 100644 --- a/modules/processing_correction.py +++ b/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: From e4cb96f6a43e015b018c162e4c0032d6365183a0 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Thu, 19 Mar 2026 04:03:20 +0000 Subject: [PATCH 3/3] fix(ui): widen latent correction slider ranges MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit brightness and sharpen: ±1.0 → ±4.0 color: 0-4.0 → 0-16.0 tint strength: ±1.0 → ±4.0 --- modules/ui_sections.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/modules/ui_sections.py b/modules/ui_sections.py index 654c7ee4b..886e392a2 100644 --- a/modules/ui_sections.py +++ b/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