| | 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) |
| |
|
