From b8ea66b262d7e89c76d5dbfbd1e05bb16a75a30b Mon Sep 17 00:00:00 2001 From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com> Date: Mon, 5 Feb 2024 10:09:24 -0600 Subject: [PATCH] Add Latent Tint --- modules/control/run.py | 4 ++-- modules/img2img.py | 4 ++-- modules/processing_class.py | 4 +++- modules/processing_correction.py | 25 +++++++++++++++++++++---- modules/txt2img.py | 4 ++-- modules/ui_control.py | 4 ++-- modules/ui_img2img.py | 4 ++-- modules/ui_sections.py | 6 +++++- modules/ui_txt2img.py | 4 ++-- scripts/xyz_grid.py | 2 ++ 10 files changed, 43 insertions(+), 18 deletions(-) diff --git a/modules/control/run.py b/modules/control/run.py index a648282dc..99b03d879 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, cfg_end, full_quality, restore_faces, tiling, - hdr_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry, + hdr_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry, hdr_color_picker, hdr_tint_ratio, 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, @@ -89,7 +89,7 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_ inpainting_mask_invert = 1 if masking.opts.invert else 0, inpainting_fill = 1, hdr_mode=hdr_mode, hdr_brightness=hdr_brightness, hdr_color=hdr_color, hdr_sharpen=hdr_sharpen, hdr_clamp=hdr_clamp, - hdr_boundary=hdr_boundary, hdr_threshold=hdr_threshold, hdr_maximize=hdr_maximize, hdr_max_center=hdr_max_center, hdr_max_boundry=hdr_max_boundry, + hdr_boundary=hdr_boundary, hdr_threshold=hdr_threshold, hdr_maximize=hdr_maximize, hdr_max_center=hdr_max_center, hdr_max_boundry=hdr_max_boundry, hdr_color_picker=hdr_color_picker, hdr_tint_ratio=hdr_tint_ratio, outpath_samples=shared.opts.outdir_samples or shared.opts.outdir_control_samples, outpath_grids=shared.opts.outdir_grids or shared.opts.outdir_control_grids, ) diff --git a/modules/img2img.py b/modules/img2img.py index 3a509b6ca..6d1805525 100644 --- a/modules/img2img.py +++ b/modules/img2img.py @@ -135,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_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry, + hdr_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry, hdr_color_picker, hdr_tint_ratio, override_settings_texts, *args): # pylint: disable=unused-argument @@ -240,7 +240,7 @@ def img2img(id_task: str, mode: int, inpaint_full_res_padding=inpaint_full_res_padding, inpainting_mask_invert=inpainting_mask_invert, hdr_mode=hdr_mode, hdr_brightness=hdr_brightness, hdr_color=hdr_color, hdr_sharpen=hdr_sharpen, hdr_clamp=hdr_clamp, - hdr_boundary=hdr_boundary, hdr_threshold=hdr_threshold, hdr_maximize=hdr_maximize, hdr_max_center=hdr_max_center, hdr_max_boundry=hdr_max_boundry, + hdr_boundary=hdr_boundary, hdr_threshold=hdr_threshold, hdr_maximize=hdr_maximize, hdr_max_center=hdr_max_center, hdr_max_boundry=hdr_max_boundry, hdr_color_picker=hdr_color_picker, hdr_tint_ratio=hdr_tint_ratio, override_settings=override_settings, ) if selected_scale_tab == 1 and resize_mode != 0: diff --git a/modules/processing_class.py b/modules/processing_class.py index 22fc27a79..a4c3cfb63 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, cfg_end: float = 1, resize_mode: int = 0, resize_name: str = 'None', scale_by: float = 0, selected_scale_tab: int = 0, hdr_mode: int = 0, hdr_brightness: float = 0, hdr_color: float = 0, hdr_sharpen: float = 0, hdr_clamp: bool = False, hdr_boundary: float = 4.0, hdr_threshold: float = 0.95, 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_mode: int = 0, hdr_brightness: float = 0, hdr_color: float = 0, hdr_sharpen: float = 0, hdr_clamp: bool = False, hdr_boundary: float = 4.0, hdr_threshold: float = 0.95, hdr_maximize: bool = False, hdr_max_center: float = 0.6, hdr_max_boundry: float = 1.0, hdr_color_picker: str = None, hdr_tint_ratio: float = 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 @@ -125,6 +125,8 @@ class StableDiffusionProcessing: self.hdr_maximize=hdr_maximize self.hdr_max_center=hdr_max_center self.hdr_max_boundry=hdr_max_boundry + self.hdr_color_picker=hdr_color_picker + self.hdr_tint_ratio=hdr_tint_ratio # globals self.scheduled_prompt: bool = False self.prompt_embeds = [] diff --git a/modules/processing_correction.py b/modules/processing_correction.py index 171d40358..9593856cc 100644 --- a/modules/processing_correction.py +++ b/modules/processing_correction.py @@ -5,8 +5,7 @@ https://huggingface.co/blog/TimothyAlexisVass/explaining-the-sdxl-latent-space import os import torch -from modules import shared - +from modules import shared, sd_vae_taesd, devices debug = shared.log.trace if os.environ.get('SD_HDR_DEBUG', None) is not None else lambda *args, **kwargs: None debug('Trace: HDR') @@ -67,11 +66,26 @@ def maximize_tensor(tensor, boundary=1.0): return tensor +def get_color(colorstr): + rgb = torch.tensor(tuple(int(colorstr.lstrip('#')[i:i + 2], 16) for i in (0, 2, 4))).to(dtype=torch.float32) + rgb = (rgb / 255).unsqueeze(-1).unsqueeze(-1).repeat(1, 64, 64).to(dtype=devices.dtype, device=devices.device) + color = sd_vae_taesd.encode(rgb).squeeze(0)[0:3, 5, 5] + return color + + +def color_adjust(tensor, colorstr, ratio): + color = get_color(colorstr) + for i in range(3): + print(color[i]) + tensor[i] = center_tensor(tensor[i], full_shift=1, offset=color[i]*(ratio/3)) + return tensor + + 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 500 < timestep < 800 and (p.hdr_brightness != 0 or p.hdr_color != 0): + if 500 < timestep < 800 and (p.hdr_brightness != 0 or p.hdr_color != 0 or p.hdr_tint_ratio != 0): p.extra_generation_params["HDR center"] = f'{p.hdr_color}/{p.hdr_brightness}' if p.hdr_brightness != 0: latent[0:1] = center_tensor(latent[0:1], full_shift=float(p.hdr_mode), offset=2*p.hdr_brightness) # Brightness @@ -81,6 +95,9 @@ def correction(p, timestep, latent): latent[1:] = center_tensor(latent[1:], channel_shift=p.hdr_color, full_shift=float(p.hdr_mode)) # Color p.extra_generation_params["HDR color"] = f'{p.hdr_color}' p.hdr_color = 0 + if p.hdr_tint_ratio != 0: + latent = color_adjust(latent, p.hdr_color_picker, p.hdr_tint_ratio) + p.hdr_tint_ratio = 0 if timestep < 350 and p.hdr_sharpen != 0: p.extra_generation_params["HDR sharpen"] = f'{p.hdr_sharpen}' per_step_ratio = 2 ** (timestep / 250) * p.hdr_sharpen / 16 @@ -95,7 +112,7 @@ def correction(p, timestep, latent): def correction_callback(p, timestep, kwargs): - if not any([p.hdr_clamp, p.hdr_mode, p.hdr_maximize, p.hdr_sharpen, p.hdr_color, p.hdr_brightness]): + 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]): return kwargs latents = kwargs["latents"] debug('') diff --git a/modules/txt2img.py b/modules/txt2img.py index f48f690f7..26b51c6f0 100644 --- a/modules/txt2img.py +++ b/modules/txt2img.py @@ -20,7 +20,7 @@ def txt2img(id_task, 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_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry, + hdr_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry, hdr_color_picker, hdr_tint_ratio, override_settings_texts, *args): @@ -78,7 +78,7 @@ def txt2img(id_task, refiner_prompt=refiner_prompt, refiner_negative=refiner_negative, hdr_mode=hdr_mode, hdr_brightness=hdr_brightness, hdr_color=hdr_color, hdr_sharpen=hdr_sharpen, hdr_clamp=hdr_clamp, - hdr_boundary=hdr_boundary, hdr_threshold=hdr_threshold, hdr_maximize=hdr_maximize, hdr_max_center=hdr_max_center, hdr_max_boundry=hdr_max_boundry, + hdr_boundary=hdr_boundary, hdr_threshold=hdr_threshold, hdr_maximize=hdr_maximize, hdr_max_center=hdr_max_center, hdr_max_boundry=hdr_max_boundry, hdr_color_picker=hdr_color_picker, hdr_tint_ratio=hdr_tint_ratio, override_settings=override_settings, ) p.scripts = scripts.scripts_txt2img diff --git a/modules/ui_control.py b/modules/ui_control.py index 17e5282e0..03b66e9d2 100644 --- a/modules/ui_control.py +++ b/modules/ui_control.py @@ -104,7 +104,7 @@ 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, cfg_end, full_quality, restore_faces, tiling= ui_sections.create_advanced_inputs('control') - hdr_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry = ui_sections.create_correction_inputs('control') + hdr_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry, hdr_color_picker, hdr_tint_ratio, = ui_sections.create_correction_inputs('control') with gr.Accordion(open=False, label="Video", elem_id="control_video", elem_classes=["small-accordion"]): with gr.Row(): @@ -489,7 +489,7 @@ def create_ui(_blocks: gr.Blocks=None): 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, cfg_end, full_quality, restore_faces, tiling, - hdr_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry, + hdr_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry, hdr_color_picker, hdr_tint_ratio, 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 511e8f935..404c578dc 100644 --- a/modules/ui_img2img.py +++ b/modules/ui_img2img.py @@ -131,7 +131,7 @@ def create_ui(): 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, cfg_end, full_quality, restore_faces, tiling = ui_sections.create_advanced_inputs('img2img') - hdr_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry = ui_sections.create_correction_inputs('img2img') + hdr_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry, hdr_color_picker, hdr_tint_ratio, = ui_sections.create_correction_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: @@ -191,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_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry, + hdr_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry, hdr_color_picker, hdr_tint_ratio, override_settings, ] img2img_dict = dict( diff --git a/modules/ui_sections.py b/modules/ui_sections.py index 10433b27b..8ae7dca17 100644 --- a/modules/ui_sections.py +++ b/modules/ui_sections.py @@ -136,7 +136,11 @@ def create_correction_inputs(tab): 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 hdr_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry + with gr.Row(): + hdr_tint_ratio = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.0, label='Tint', elem_id=f"{tab}_hdr_tint_ratio") + hdr_color_picker = gr.ColorPicker(label="Color Replace", show_label=False, value=None, elem_id=f"{tab}_hdr_color_picker") + return hdr_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry, hdr_color_picker, hdr_tint_ratio, + 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 c978a33e8..f97c97012 100644 --- a/modules/ui_txt2img.py +++ b/modules/ui_txt2img.py @@ -46,7 +46,7 @@ def create_ui(): 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, cfg_end, full_quality, restore_faces, tiling = ui_sections.create_advanced_inputs('txt2img') - hdr_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry = ui_sections.create_correction_inputs('txt2img') + hdr_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry, hdr_color_picker, hdr_tint_ratio, = ui_sections.create_correction_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') @@ -80,7 +80,7 @@ def create_ui(): 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_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry, + hdr_mode, hdr_brightness, hdr_color, hdr_sharpen, hdr_clamp, hdr_boundary, hdr_threshold, hdr_maximize, hdr_max_center, hdr_max_boundry, hdr_color_picker, hdr_tint_ratio, override_settings, ] txt2img_dict = dict( diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py index ac5458d90..a76a668e4 100644 --- a/scripts/xyz_grid.py +++ b/scripts/xyz_grid.py @@ -265,6 +265,8 @@ axis_options = [ AxisOption("[HDR] Clamp threshold", float, apply_field("hdr_threshold")), AxisOption("[HDR] Maximize center shift", float, apply_field("hdr_max_center")), AxisOption("[HDR] Maximize boundary", float, apply_field("hdr_max_boundry")), + AxisOption("[HDR] Tint Color Hex", str, apply_field("hdr_color_picker")), + AxisOption("[HDR] Tint Ratio", float, apply_field("hdr_tint_ratio")), 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')), AxisOption("[FreeU] 1st stage backbone factor", float, apply_setting('freeu_b1')),