| from modules.sd_samplers_kdiffusion import KDiffusionSampler |
| from modules import shared, scripts, script_callbacks |
| import gradio as gr |
|
|
| from scripts.cg_xyz import xyz_support |
|
|
| VERSION = "v1.1.0" |
|
|
| DYNAMIC_RANGE = [3.25, 2.5, 2.5, 2.5] |
|
|
| Default_LUTs = {"C": 0.01, "M": 0.5, "Y": -0.13, "K": 0} |
|
|
|
|
| def normalize_tensor(x, r): |
| X = x.detach().clone() |
|
|
| ratio = r / max(abs(float(X.min())), abs(float(X.max()))) |
| X *= max(ratio, 0.99) |
|
|
| return X |
|
|
|
|
| original_callback = KDiffusionSampler.callback_state |
|
|
|
|
| def center_callback(self, d): |
| if not self.diffcg_enable or getattr(self.p, "_ad_inner", False): |
| return original_callback(self, d) |
|
|
| X = d["denoised"].detach().clone() |
| batchSize = X.size(0) |
| channels = len(self.LUTs) |
|
|
| for b in range(batchSize): |
| for c in range(channels): |
|
|
| if self.diffcg_recenter_strength > 0.0: |
| d["denoised"][b][c] += ( |
| self.LUTs[c] - X[b][c].mean() |
| ) * self.diffcg_recenter_strength |
|
|
| if self.diffcg_normalize and (d["i"] + 1) > self.diffcg_last_step // 2: |
| d["denoised"][b][c] = normalize_tensor(X[b][c], DYNAMIC_RANGE[c]) |
|
|
| return original_callback(self, d) |
|
|
|
|
| KDiffusionSampler.callback_state = center_callback |
|
|
|
|
| |
| ac = getattr(shared.opts, "always_center", "None") |
| an = getattr(shared.opts, "always_normalize", "None") |
| def_sd = getattr(shared.opts, "default_arch", "1.5") |
| adv_opt = getattr(shared.opts, "show_center_opt", False) |
|
|
| c_t2i = ac in ("txt2img", "Both") |
| c_i2i = ac in ("img2img", "Both") |
| n_t2i = an in ("txt2img", "Both") |
| n_i2i = an in ("img2img", "Both") |
|
|
|
|
| class DiffusionCG(scripts.Script): |
| def __init__(self): |
| self.xyzCache = {} |
| xyz_support(self.xyzCache) |
|
|
| def title(self): |
| return "DiffusionCG" |
|
|
| def show(self, is_img2img): |
| return scripts.AlwaysVisible |
|
|
| def ui(self, is_img2img): |
| with gr.Accordion(f"Diffusion CG {VERSION}", open=False): |
| with gr.Row(): |
| enableG = gr.Checkbox( |
| label="Enable (Global)", |
| value=( |
| ((not is_img2img) and (c_t2i or n_t2i)) |
| or (is_img2img and (c_i2i or n_i2i)) |
| ), |
| ) |
| sd_ver = gr.Radio( |
| ["1.5", "XL"], value=def_sd, label="Stable Diffusion Version" |
| ) |
|
|
| with gr.Row(): |
| with gr.Group(): |
| gr.Markdown('<h3 align="center">Recenter</h3>') |
|
|
| if not is_img2img: |
| v = 1.0 if c_t2i else 0.0 |
| else: |
| v = 1.0 if c_i2i else 0.0 |
|
|
| rc_str = gr.Slider( |
| label="Effect Strength", |
| minimum=0.0, |
| maximum=1.0, |
| step=0.2, |
| value=v, |
| ) |
|
|
| with gr.Group(): |
| gr.Markdown('<h3 align="center">Normalization</h3>') |
| enableN = gr.Checkbox( |
| label="Activate", |
| value=(((not is_img2img) and n_t2i) or (is_img2img and n_i2i)), |
| ) |
|
|
| with gr.Accordion("Recenter Settings", visible=adv_opt, open=False): |
| with gr.Group(visible=(def_sd == "1.5")) as setting15: |
| C = gr.Slider( |
| label="C", |
| minimum=-1.00, |
| maximum=1.00, |
| step=0.01, |
| value=Default_LUTs["C"], |
| ) |
| M = gr.Slider( |
| label="M", |
| minimum=-1.00, |
| maximum=1.00, |
| step=0.01, |
| value=Default_LUTs["M"], |
| ) |
| Y = gr.Slider( |
| label="Y", |
| minimum=-1.00, |
| maximum=1.00, |
| step=0.01, |
| value=Default_LUTs["Y"], |
| ) |
| K = gr.Slider( |
| label="K", |
| minimum=-1.00, |
| maximum=1.00, |
| step=0.01, |
| value=Default_LUTs["K"], |
| ) |
|
|
| with gr.Group(visible=(def_sd == "XL")) as settingXL: |
| L = gr.Slider( |
| label="L", minimum=-1.00, maximum=1.00, step=0.01, value=0.0 |
| ) |
| a = gr.Slider( |
| label="a", minimum=-1.00, maximum=1.00, step=0.01, value=0.0 |
| ) |
| b = gr.Slider( |
| label="b", minimum=-1.00, maximum=1.00, step=0.01, value=0.0 |
| ) |
|
|
| def on_radio_change(choice): |
| if choice == "1.5": |
| return [ |
| gr.Group.update(visible=True), |
| gr.Group.update(visible=False), |
| ] |
| else: |
| return [ |
| gr.Group.update(visible=False), |
| gr.Group.update(visible=True), |
| ] |
|
|
| sd_ver.change(on_radio_change, sd_ver, [setting15, settingXL]) |
|
|
| self.paste_field_names = [ |
| (rc_str, "ReCenter Str"), |
| (enableN, "Normalization"), |
| (sd_ver, "SD_ver"), |
| ] |
| self.infotext_fields = [ |
| (rc_str, "ReCenter Str"), |
| (enableN, "Normalization"), |
| (sd_ver, "SD_ver"), |
| ] |
|
|
| if adv_opt: |
| self.paste_field_names += [ |
| ( |
| C, |
| lambda d: ( |
| float(d["LUTs"].strip("[]").split(",")[0]) |
| if len(d.get("LUTs", "").split(",")) == 4 |
| else gr.update() |
| ), |
| ), |
| ( |
| M, |
| lambda d: ( |
| float(d["LUTs"].strip("[]").split(",")[1]) |
| if len(d.get("LUTs", "").split(",")) == 4 |
| else gr.update() |
| ), |
| ), |
| ( |
| Y, |
| lambda d: ( |
| float(d["LUTs"].strip("[]").split(",")[2]) |
| if len(d.get("LUTs", "").split(",")) == 4 |
| else gr.update() |
| ), |
| ), |
| ( |
| K, |
| lambda d: ( |
| float(d["LUTs"].strip("[]").split(",")[3]) |
| if len(d.get("LUTs", "").split(",")) == 4 |
| else gr.update() |
| ), |
| ), |
| ( |
| L, |
| lambda d: ( |
| float(d["LUTs"].strip("[]").split(",")[0]) |
| if len(d.get("LUTs", "").split(",")) == 3 |
| else gr.update() |
| ), |
| ), |
| ( |
| a, |
| lambda d: ( |
| float(d["LUTs"].strip("[]").split(",")[1]) |
| if len(d.get("LUTs", "").split(",")) == 3 |
| else gr.update() |
| ), |
| ), |
| ( |
| b, |
| lambda d: ( |
| float(d["LUTs"].strip("[]").split(",")[2]) |
| if len(d.get("LUTs", "").split(",")) == 3 |
| else gr.update() |
| ), |
| ), |
| ] |
| self.infotext_fields += [ |
| ( |
| C, |
| lambda d: ( |
| float(d["LUTs"].strip("[]").split(",")[0]) |
| if len(d.get("LUTs", "").split(",")) == 4 |
| else gr.update() |
| ), |
| ), |
| ( |
| M, |
| lambda d: ( |
| float(d["LUTs"].strip("[]").split(",")[1]) |
| if len(d.get("LUTs", "").split(",")) == 4 |
| else gr.update() |
| ), |
| ), |
| ( |
| Y, |
| lambda d: ( |
| float(d["LUTs"].strip("[]").split(",")[2]) |
| if len(d.get("LUTs", "").split(",")) == 4 |
| else gr.update() |
| ), |
| ), |
| ( |
| K, |
| lambda d: ( |
| float(d["LUTs"].strip("[]").split(",")[3]) |
| if len(d.get("LUTs", "").split(",")) == 4 |
| else gr.update() |
| ), |
| ), |
| ( |
| L, |
| lambda d: ( |
| float(d["LUTs"].strip("[]").split(",")[0]) |
| if len(d.get("LUTs", "").split(",")) == 3 |
| else gr.update() |
| ), |
| ), |
| ( |
| a, |
| lambda d: ( |
| float(d["LUTs"].strip("[]").split(",")[1]) |
| if len(d.get("LUTs", "").split(",")) == 3 |
| else gr.update() |
| ), |
| ), |
| ( |
| b, |
| lambda d: ( |
| float(d["LUTs"].strip("[]").split(",")[2]) |
| if len(d.get("LUTs", "").split(",")) == 3 |
| else gr.update() |
| ), |
| ), |
| ] |
|
|
| for comp in [enableG, sd_ver, rc_str, enableN, C, M, Y, K, L, a, b]: |
| comp.do_not_save_to_config = True |
|
|
| return [enableG, sd_ver, rc_str, enableN, C, M, Y, K, L, a, b] |
|
|
| def before_hr(self, p, *args): |
| KDiffusionSampler.diffcg_normalize = False |
|
|
| def process( |
| self, |
| p, |
| enableG: bool, |
| sd_ver: str, |
| rc_str: float, |
| enableN: bool, |
| C: float, |
| M: float, |
| Y: float, |
| K: float, |
| L: float, |
| a: float, |
| b: float, |
| ): |
|
|
| if "enableG" in self.xyzCache.keys(): |
| enableG = self.xyzCache["enableG"].lower().strip() == "true" |
| del self.xyzCache["enableG"] |
|
|
| KDiffusionSampler.diffcg_enable = enableG |
| if not enableG: |
| if len(self.xyzCache.keys()) > 0: |
| print("\n[Diff. CG] X [X/Y/Z Plot] Extension is not Enabled!\n") |
| self.xyzCache.clear() |
| return p |
|
|
| if "rc_str" in self.xyzCache.keys(): |
| rc_str = float(self.xyzCache["rc_str"]) |
| if "enableN" in self.xyzCache.keys(): |
| enableN = self.xyzCache["enableN"].lower().strip() == "true" |
|
|
| if adv_opt: |
| C = self.xyzCache.get("C", C) |
| M = self.xyzCache.get("M", M) |
| Y = self.xyzCache.get("Y", Y) |
| K = self.xyzCache.get("K", K) |
| L = self.xyzCache.get("L", L) |
| a = self.xyzCache.get("a", a) |
| b = self.xyzCache.get("b", b) |
|
|
| if sd_ver == "1.5": |
| KDiffusionSampler.LUTs = [-K, -M, C, Y] |
| else: |
| KDiffusionSampler.LUTs = [L, -a, b] |
|
|
| KDiffusionSampler.diffcg_recenter_strength = rc_str |
| KDiffusionSampler.diffcg_normalize = enableN |
|
|
| if ( |
| not hasattr(p, "enable_hr") |
| and not shared.opts.img2img_fix_steps |
| and getattr(p, "denoising_strength", 1.0) < 1.0 |
| ): |
| KDiffusionSampler.diffcg_last_step = int(p.steps * p.denoising_strength) + 1 |
| else: |
| KDiffusionSampler.diffcg_last_step = p.steps |
|
|
| p.extra_generation_params.update( |
| { |
| "ReCenter Str": rc_str, |
| "Normalization": enableN, |
| "SD_ver": sd_ver, |
| } |
| ) |
|
|
| if adv_opt: |
| if def_sd == "1.5": |
| p.extra_generation_params["LUTs"] = f"[{C}, {M}, {Y}, {K}]" |
| else: |
| p.extra_generation_params["LUTs"] = f"[{L}, {a}, {b}]" |
|
|
| self.xyzCache.clear() |
|
|
|
|
| def restore_callback(): |
| KDiffusionSampler.callback_state = original_callback |
|
|
|
|
| script_callbacks.on_script_unloaded(restore_callback) |
|
|
