├── .gitignore ├── CHANGELOG.md ├── LICENSE ├── README.md ├── examples ├── nml_off.jpg ├── nml_on.jpg ├── rc_off.jpg ├── rc_on.jpg ├── xl_off.jpg └── xl_on.jpg ├── lib_cg ├── __init__.py ├── adv.py ├── params.py ├── settings.py └── xyz.py └── scripts └── diffusion_cg.py /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | ### v1.2.0 - 2024 Aug.30 2 | - Major **Rewrite** 3 | 4 | ### v1.1.1 - 2024 Aug.30 5 | - Add Settings **Section** 6 | 7 | ### v1.1.0 - 2024 May.06 8 | - Add **X/Y/Z Plot** Support 9 | 10 | ### v1.0.0 - 2024 Mar.07 11 | - Extension "**Released**" 12 | 13 | ### v0.5.0 - 2024 Jan.09 14 | - Add **Infotext** 15 | 16 | ### v0.4.2 - 2023 Dec.19 17 | - Optimization via `.detach()` 18 | 19 | ### v0.4.1 - 2023 Dec.11 20 | - Update **Normalization** Logic 21 | 22 | ### v0.4.0 - 2023 Dec.08 23 | - Implement Color **Parameters** 24 | 25 | ### v0.3.0 - 2023 Dec.01 26 | - Add **Default** Settings 27 | 28 | ### v0.2.3 - 2023 Dec.01 29 | - **Clarify** README 30 | 31 | ### v0.2.2 - 2023 Nov.25 32 | - No longer process if a **Mask** exists 33 | 34 | ### v0.2.1 - 2023 Nov.23 35 | - Add **Effect Strength** Settings 36 | 37 | ### v0.2.0 - 2023 Nov.21 38 | - Better **SDXL** Support 39 | 40 | ### v0.1.3 - 2023 Nov.20 41 | - Improved **Compatibility** 42 | - 43 | ### v0.1.2 - 2023 Nov.20 44 | - Improved(?) **Normalization** 45 | 46 | ### v0.1.1 - 2023 Nov.20 47 | - Improved(?) **Normalization** 48 | 49 | ### v0.1.0 - 2023 Nov.16 50 | - Public **Alpha** Released~ 51 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 Haoming 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # SD Webui Diffusion Color Grading 2 | This is an Extension for the [Automatic1111 Webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui), which performs *Color Grading* during the generation, producing a more **neutral** and **balanced**, but also **vibrant** and **contrasty** color. 3 | 4 | > This is the fruition of the joint research between [TimothyAlexisVass](https://github.com/TimothyAlexisVass) with their findings, and me with my experience in developing [Vectorscope CC](https://github.com/Haoming02/sd-webui-vectorscope-cc) 5 | 6 | **Note:** This Extension is disabled during [ADetailer](https://github.com/Bing-su/adetailer) phase to prevent inconsistent colors 7 | 8 | ## Features 9 | 10 | This Extension comes with two main effects, **Recenter** and **Normalization**: 11 | 12 | #### Recenter 13 | 14 |
Abstract
15 | 16 | TimothyAlexisVass discovered that, the value of the latent noise Tensor often starts off-centered, and the mean of each channel tends to drift away from `0`. Therefore, I wrote a function to guide the mean back to `0`. 17 | 18 |
Effects
19 | 20 | When you enable the feature, the output images will not have a biased color tint, and all colors will distribute more evenly; Additionally, the brightness will be adjusted so that bright areas are not overblown and dark areas are not clipped, producing an effect similar to HDR photos. 21 | 22 |
Samples
23 | 24 |

25 | 26 | 27 |
Off | On
28 |

29 | 30 | #### Normalization 31 | 32 |
Abstract
33 | 34 | By encoding images back into latent noise using VAE, TimothyAlexisVass discovered that the resulting values usually fall within a certain range, and thus theorized that if the final latent noise has a smaller value range than normal, then some precision is essentailly wasted. This gave me an idea to write a function that make the latent noise utilize the full depth. 35 | 36 |
Effects
37 | 38 | When you enable the feature, the latent noise will attempt to span across the full value range if possible, before getting decoded by the VAE. As a result, bright areas will get brighter and dark areas will get darker, and additional details may also be introduced in these areas. 39 | 40 |

41 | 42 | 43 |
Off | On
44 |

45 | 46 | > Both features increase the image file size when enabled, suggesting that they "contain more informations" 47 | 48 | #### Misc. 49 | 50 | - You can enable both features at the same time to generate some stunning images 51 | - This Extension supports both `SD 1.5` and `SDXL` checkpoints 52 | 53 |

54 | 55 | 56 |
Off | On
57 |

58 | 59 | ## Settings 60 | In the `Diffusion CG` section under the Stable Diffusion category in the **Settings** tab, you can make either feature default to `enable`, as well as setting the Stable Diffusion architecture to start with. 61 | 62 |
63 | 64 |
65 | Structures of Stable Diffusion 66 | 67 | The `Tensor` of the latent noise has a dimention of `[batch, 4, height / 8, width / 8]`. 68 | 69 | - For **SD 1.5:** From my trial and error when developing [Vectorscope CC](https://github.com/Haoming02/sd-webui-vectorscope-cc), each of the 4 channels essentially represents the `-K`, `-M`, `C`, `Y` color for the **CMYK** color model. 70 | 71 | - For **SDXL:** According to TimothyAlexisVass's [Blogpost](https://huggingface.co/blog/TimothyAlexisVass/explaining-the-sdxl-latent-space), the first 3 channels represent the `Y'`, `-Cr`, `-Cb` color for the **Y'CbCr** color model, while the 4th channel is the pattern/structure. 72 | 73 |
74 | -------------------------------------------------------------------------------- /examples/nml_off.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Haoming02/sd-webui-diffusion-cg/43a9fcce8b350afd0442f604541eb580a514849d/examples/nml_off.jpg -------------------------------------------------------------------------------- /examples/nml_on.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Haoming02/sd-webui-diffusion-cg/43a9fcce8b350afd0442f604541eb580a514849d/examples/nml_on.jpg -------------------------------------------------------------------------------- /examples/rc_off.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Haoming02/sd-webui-diffusion-cg/43a9fcce8b350afd0442f604541eb580a514849d/examples/rc_off.jpg -------------------------------------------------------------------------------- /examples/rc_on.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Haoming02/sd-webui-diffusion-cg/43a9fcce8b350afd0442f604541eb580a514849d/examples/rc_on.jpg -------------------------------------------------------------------------------- /examples/xl_off.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Haoming02/sd-webui-diffusion-cg/43a9fcce8b350afd0442f604541eb580a514849d/examples/xl_off.jpg -------------------------------------------------------------------------------- /examples/xl_on.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Haoming02/sd-webui-diffusion-cg/43a9fcce8b350afd0442f604541eb580a514849d/examples/xl_on.jpg -------------------------------------------------------------------------------- /lib_cg/__init__.py: -------------------------------------------------------------------------------- 1 | from modules import shared 2 | from .settings import * 3 | 4 | SCALE_FACTOR: dict = { 5 | "1.5": 0.18215, # https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/v1.10.1/configs/v1-inference.yaml#L17 6 | "XL": 0.13025, # https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/v1.10.1/configs/sd_xl_inpaint.yaml#L4 7 | } 8 | 9 | # ["None", "txt2img", "img2img", "Both"] 10 | ac: str = getattr(shared.opts, "always_center", "None") 11 | an: str = getattr(shared.opts, "always_normalize", "None") 12 | 13 | c_t2i: bool = ac in ("txt2img", "Both") 14 | c_i2i: bool = ac in ("img2img", "Both") 15 | n_t2i: bool = an in ("txt2img", "Both") 16 | n_i2i: bool = an in ("img2img", "Both") 17 | 18 | default_sd: str = getattr(shared.opts, "default_arch", "1.5") 19 | -------------------------------------------------------------------------------- /lib_cg/adv.py: -------------------------------------------------------------------------------- 1 | import gradio as gr 2 | 3 | COLOR_DESC: dict[str, str] = { 4 | "C": "Red | Cyan", 5 | "M": "Green | Magenta", 6 | "Y": "Blue | Yellow", 7 | "K": "Bright | Dark", 8 | "Y'": "Dark | Bright", 9 | "Cb": None, 10 | "Cr": None, 11 | } 12 | 13 | 14 | def advanced_settings(default_sd: str, sd_ver: gr.Radio) -> list[gr.Slider]: 15 | 16 | color_channels: list[gr.Slider] = [] 17 | 18 | with gr.Group(visible=(default_sd == "1.5")) as setting15: 19 | 20 | for ch in ("C", "M", "Y", "K"): 21 | color_channels.append( 22 | gr.Slider( 23 | label=ch, 24 | info=COLOR_DESC[ch], 25 | minimum=-1.00, 26 | maximum=1.00, 27 | step=0.05, 28 | value=0.0, 29 | ) 30 | ) 31 | 32 | with gr.Group(visible=(default_sd == "XL")) as settingXL: 33 | 34 | for ch in ("Y'", "Cb", "Cr"): 35 | color_channels.append( 36 | gr.Slider( 37 | label=ch, 38 | info=COLOR_DESC[ch], 39 | minimum=-1.00, 40 | maximum=1.00, 41 | step=0.05, 42 | value=0.0, 43 | ) 44 | ) 45 | 46 | def on_arch_change(choice: str): 47 | return [ 48 | gr.update(visible=(choice == "1.5")), 49 | gr.update(visible=(choice == "XL")), 50 | ] 51 | 52 | sd_ver.change(on_arch_change, inputs=[sd_ver], outputs=[setting15, settingXL]) 53 | 54 | return color_channels 55 | -------------------------------------------------------------------------------- /lib_cg/params.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass 2 | 3 | 4 | @dataclass 5 | class DiffusionCGParams: 6 | 7 | total_step: int 8 | 9 | enable: bool 10 | sd_ver: str 11 | 12 | rc_str: float 13 | LUTs: list[float] 14 | 15 | normalization: bool 16 | dynamic_range: float 17 | 18 | def __bool__(self): 19 | return self.enable 20 | -------------------------------------------------------------------------------- /lib_cg/settings.py: -------------------------------------------------------------------------------- 1 | from modules.script_callbacks import on_ui_settings 2 | from modules.shared import OptionInfo, opts 3 | import gradio as gr 4 | 5 | section = ("diffusion_cg", "Diffusion CG") 6 | tabs = ("None", "txt2img", "img2img", "Both") 7 | 8 | 9 | def on_settings(): 10 | opts.add_option( 11 | "always_center", 12 | OptionInfo( 13 | "None", 14 | "Always Perform Recenter on:", 15 | component=gr.Radio, 16 | component_args={"choices": tabs}, 17 | section=section, 18 | category_id="sd", 19 | ).needs_reload_ui(), 20 | ) 21 | 22 | opts.add_option( 23 | "always_normalize", 24 | OptionInfo( 25 | "None", 26 | "Always Perform Normalization on:", 27 | component=gr.Radio, 28 | component_args={"choices": tabs}, 29 | section=section, 30 | category_id="sd", 31 | ).needs_reload_ui(), 32 | ) 33 | 34 | opts.add_option( 35 | "default_arch", 36 | OptionInfo( 37 | "1.5", 38 | "Default Stable Diffusion Version", 39 | component=gr.Radio, 40 | component_args={"choices": ("1.5", "XL")}, 41 | section=section, 42 | category_id="sd", 43 | ).needs_reload_ui(), 44 | ) 45 | 46 | 47 | on_ui_settings(on_settings) 48 | -------------------------------------------------------------------------------- /lib_cg/xyz.py: -------------------------------------------------------------------------------- 1 | from modules import scripts 2 | 3 | 4 | def grid_reference(): 5 | for data in scripts.scripts_data: 6 | if data.script_class.__module__ in ( 7 | "scripts.xyz_grid", 8 | "xyz_grid.py", 9 | ) and hasattr(data, "module"): 10 | return data.module 11 | 12 | raise SystemError("Could not find X/Y/Z Plot...") 13 | 14 | 15 | def xyz_support(cache: dict): 16 | 17 | def apply_field(field): 18 | def _(p, x, xs): 19 | cache.update({field: x}) 20 | 21 | return _ 22 | 23 | def choices_bool(): 24 | return ["False", "True"] 25 | 26 | xyz_grid = grid_reference() 27 | 28 | extra_axis_options = [ 29 | xyz_grid.AxisOption( 30 | "[Diff. CG] Enable", 31 | str, 32 | apply_field("enable"), 33 | choices=choices_bool, 34 | ), 35 | xyz_grid.AxisOption( 36 | "[Diff. CG] Recenter", 37 | float, 38 | apply_field("rc_str"), 39 | ), 40 | xyz_grid.AxisOption( 41 | "[Diff. CG] Normalization", 42 | str, 43 | apply_field("normalization"), 44 | choices=choices_bool, 45 | ), 46 | ] 47 | 48 | extra_axis_options.extend( 49 | [ 50 | xyz_grid.AxisOption("[Diff.CG] C", float, apply_field("C")), 51 | xyz_grid.AxisOption("[Diff.CG] M", float, apply_field("M")), 52 | xyz_grid.AxisOption("[Diff.CG] Y", float, apply_field("Y")), 53 | xyz_grid.AxisOption("[Diff.CG] K", float, apply_field("K")), 54 | xyz_grid.AxisOption("[Diff.CG] Y'", float, apply_field("Lu")), 55 | xyz_grid.AxisOption("[Diff.CG] Cb", float, apply_field("Cb")), 56 | xyz_grid.AxisOption("[Diff.CG] Cr", float, apply_field("Cr")), 57 | ] 58 | ) 59 | 60 | xyz_grid.axis_options.extend(extra_axis_options) 61 | -------------------------------------------------------------------------------- /scripts/diffusion_cg.py: -------------------------------------------------------------------------------- 1 | from modules.sd_samplers_kdiffusion import KDiffusionSampler 2 | from modules.script_callbacks import on_script_unloaded 3 | from modules.ui_components import InputAccordion 4 | from modules import scripts 5 | 6 | from functools import wraps 7 | import gradio as gr 8 | import torch 9 | 10 | from lib_cg import SCALE_FACTOR, c_t2i, c_i2i, n_t2i, n_i2i, default_sd 11 | from lib_cg.params import DiffusionCGParams 12 | from lib_cg.adv import advanced_settings 13 | from lib_cg.xyz import xyz_support 14 | 15 | VERSION = "1.2.0" 16 | 17 | 18 | @torch.inference_mode() 19 | def normalize_tensor(x: torch.Tensor, r: float) -> torch.Tensor: 20 | ratio: float = r / max(abs(float(x.min())), abs(float(x.max()))) 21 | x *= max(ratio, 1.0) 22 | return x 23 | 24 | 25 | original_callback = KDiffusionSampler.callback_state 26 | 27 | 28 | @torch.inference_mode() 29 | @wraps(original_callback) 30 | def center_callback(self, d): 31 | if getattr(self.p, "_ad_inner", False): 32 | return original_callback(self, d) 33 | 34 | params: DiffusionCGParams = getattr(self, "diffcg_params", None) 35 | if not params: 36 | return original_callback(self, d) 37 | 38 | X: torch.Tensor = d["denoised"].detach().clone() 39 | batchSize: int = X.size(0) 40 | channels: int = len(params.LUTs) 41 | 42 | for b in range(batchSize): 43 | for c in range(channels): 44 | 45 | if params.rc_str > 0.0: 46 | d["denoised"][b][c] += (params.LUTs[c] - X[b][c].mean()) * params.rc_str 47 | 48 | if params.normalization and (d["i"] + 1) >= (params.total_step - 1): 49 | d["denoised"][b][c] = normalize_tensor(X[b][c], params.dynamic_range) 50 | 51 | return original_callback(self, d) 52 | 53 | 54 | KDiffusionSampler.callback_state = center_callback 55 | 56 | 57 | class DiffusionCG(scripts.Script): 58 | 59 | def __init__(self): 60 | self.xyzCache = {} 61 | xyz_support(self.xyzCache) 62 | 63 | def title(self): 64 | return "Diffusion CG" 65 | 66 | def show(self, is_img2img): 67 | return scripts.AlwaysVisible 68 | 69 | def ui(self, is_img2img): 70 | 71 | with InputAccordion( 72 | label=f"{self.title()} v{VERSION}", 73 | value=( 74 | ((not is_img2img) and (c_t2i or n_t2i)) 75 | or (is_img2img and (c_i2i or n_i2i)) 76 | ), 77 | ) as enable: 78 | 79 | with gr.Row(): 80 | 81 | sd_ver = gr.Radio( 82 | ["1.5", "XL"], value=default_sd, label="Stable Diffusion Version" 83 | ) 84 | 85 | with gr.Column(): 86 | 87 | rc_str = gr.Slider( 88 | label="[Recenter]", 89 | info="Effect Strength", 90 | minimum=0.0, 91 | maximum=1.0, 92 | step=0.2, 93 | value=( 94 | 1.0 95 | if (((not is_img2img) and c_t2i) or (is_img2img and c_i2i)) 96 | else 0.0 97 | ), 98 | ) 99 | 100 | normalization = gr.Checkbox( 101 | label="[Normalization]", 102 | value=(((not is_img2img) and n_t2i) or (is_img2img and n_i2i)), 103 | ) 104 | 105 | with gr.Accordion("Recenter Settings", open=False): 106 | C, M, Y, K, L, a, b = advanced_settings(default_sd, sd_ver) 107 | 108 | comps: list = [] 109 | 110 | self.paste_field_names = [] 111 | self.infotext_fields = [ 112 | (enable, "Diffusion CG Enable"), 113 | (sd_ver, "Diffusion CG SD-Ver"), 114 | (rc_str, "Diffusion CG Recenter"), 115 | (normalization, "[Diff. CG] Normalization"), 116 | (C, "Diffusion CG C"), 117 | (M, "Diffusion CG M"), 118 | (Y, "Diffusion CG Y"), 119 | (K, "Diffusion CG K"), 120 | (L, "Diffusion CG Y'"), 121 | (a, "Diffusion CG Cb"), 122 | (b, "Diffusion CG Cr"), 123 | ] 124 | 125 | for comp, name in self.infotext_fields: 126 | comp.do_not_save_to_config = True 127 | self.paste_field_names.append(name) 128 | comps.append(comp) 129 | 130 | return comps 131 | 132 | def before_hr(self, p, *args): 133 | KDiffusionSampler.diffcg_last_step = ( 134 | getattr(p, "hr_second_pass_steps", None) or p.steps 135 | ) 136 | 137 | def process( 138 | self, 139 | p, 140 | enable: bool, 141 | sd_ver: str, 142 | rc_str: float, 143 | normalization: bool, 144 | C: float, 145 | M: float, 146 | Y: float, 147 | K: float, 148 | Lu: float, 149 | Cb: float, 150 | Cr: float, 151 | ): 152 | 153 | if "enable" in self.xyzCache.keys(): 154 | enable = self.xyzCache["enable"].lower().strip() == "true" 155 | 156 | if not enable: 157 | if len(self.xyzCache) > 0: 158 | if "enable" not in self.xyzCache.keys(): 159 | print("\n[Vec.CC] x [X/Y/Z Plot] Extension is not Enabled!\n") 160 | self.xyzCache.clear() 161 | 162 | setattr(KDiffusionSampler, "diffcg_params", None) 163 | return p 164 | 165 | rc_str = float(self.xyzCache.get("rc_str", rc_str)) 166 | 167 | if "normalization" in self.xyzCache.keys(): 168 | normalization = self.xyzCache["normalization"].lower().strip() == "true" 169 | 170 | C = float(self.xyzCache.get("C", C)) 171 | M = float(self.xyzCache.get("M", M)) 172 | Y = float(self.xyzCache.get("Y", Y)) 173 | K = float(self.xyzCache.get("K", K)) 174 | Lu = float(self.xyzCache.get("Lu", Lu)) 175 | Cb = float(self.xyzCache.get("Cb", Cb)) 176 | Cr = float(self.xyzCache.get("Cr", Cr)) 177 | 178 | params = DiffusionCGParams( 179 | getattr(p, "firstpass_steps", None) or p.steps, 180 | enable, 181 | sd_ver, 182 | rc_str, 183 | [-K, -M, C, Y] if (sd_ver == "1.5") else [Lu, -Cr, -Cb], 184 | normalization, 185 | (1.0 / SCALE_FACTOR["1.5"]) / 2.0, # XL causes Noises... 186 | ) 187 | 188 | setattr(KDiffusionSampler, "diffcg_params", params) 189 | 190 | p.extra_generation_params.update( 191 | { 192 | "Diffusion CG Enable": enable, 193 | "Diffusion CG SD-Ver": sd_ver, 194 | "Diffusion CG Recenter": rc_str, 195 | "Diffusion CG Normalization": normalization, 196 | } 197 | ) 198 | 199 | if sd_ver == "1.5": 200 | p.extra_generation_params.update( 201 | { 202 | "Diffusion CG C": C, 203 | "Diffusion CG M": M, 204 | "Diffusion CG Y": Y, 205 | "Diffusion CG K": K, 206 | } 207 | ) 208 | else: 209 | p.extra_generation_params.update( 210 | { 211 | "Diffusion CG Y'": Lu, 212 | "Diffusion CG Cb": Cb, 213 | "Diffusion CG Cr": Cr, 214 | } 215 | ) 216 | 217 | self.xyzCache.clear() 218 | 219 | 220 | def restore_callback(): 221 | KDiffusionSampler.callback_state = original_callback 222 | 223 | 224 | on_script_unloaded(restore_callback) 225 | --------------------------------------------------------------------------------