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