|
|
import math
|
|
|
import torch
|
|
|
import re
|
|
|
import gradio as gr
|
|
|
import numpy as np
|
|
|
import modules.scripts as scripts
|
|
|
import modules.images as saving
|
|
|
from modules import devices, processing, shared, sd_samplers_kdiffusion, sd_samplers_compvis, script_callbacks
|
|
|
from modules.processing import Processed
|
|
|
from modules.shared import opts, state
|
|
|
from ldm.models.diffusion import ddim
|
|
|
from PIL import Image
|
|
|
|
|
|
from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, noise_like
|
|
|
|
|
|
re_prompt_cfgr = re.compile(r"<cfg_rescale:([^>]+)>")
|
|
|
|
|
|
class Script(scripts.Script):
|
|
|
|
|
|
def __init__(self):
|
|
|
self.old_denoising = sd_samplers_kdiffusion.CFGDenoiser.combine_denoised
|
|
|
self.old_schedule = ddim.DDIMSampler.make_schedule
|
|
|
self.old_sample = ddim.DDIMSampler.p_sample_ddim
|
|
|
globals()['enable_furry_cocks'] = True
|
|
|
|
|
|
def find_module(module_names):
|
|
|
if isinstance(module_names, str):
|
|
|
module_names = [s.strip() for s in module_names.split(",")]
|
|
|
for data in scripts.scripts_data:
|
|
|
if data.script_class.__module__ in module_names and hasattr(data, "module"):
|
|
|
return data.module
|
|
|
return None
|
|
|
|
|
|
def rescale_opt(p, x, xs):
|
|
|
globals()['cfg_rescale_fi'] = x
|
|
|
globals()['enable_furry_cocks'] = False
|
|
|
|
|
|
xyz_grid = find_module("xyz_grid.py, xy_grid.py")
|
|
|
if xyz_grid:
|
|
|
extra_axis_options = [xyz_grid.AxisOption("Rescale CFG", float, rescale_opt)]
|
|
|
xyz_grid.axis_options.extend(extra_axis_options)
|
|
|
|
|
|
def title(self):
|
|
|
return "CFG Rescale Extension"
|
|
|
|
|
|
def show(self, is_img2img):
|
|
|
return scripts.AlwaysVisible
|
|
|
|
|
|
def ui(self, is_img2img):
|
|
|
with gr.Accordion("CFG Rescale", open=True, elem_id="cfg_rescale"):
|
|
|
rescale = gr.Slider(label="CFG Rescale", show_label=False, minimum=0.0, maximum=1.0, step=0.01, value=0.0)
|
|
|
with gr.Row():
|
|
|
recolor = gr.Checkbox(label="Auto Color Fix", default=False)
|
|
|
rec_strength = gr.Slider(label="Fix Strength", interactive=True, visible=False,
|
|
|
elem_id=self.elem_id("rec_strength"), minimum=0.1, maximum=10.0, step=0.1,
|
|
|
value=1.0)
|
|
|
show_original = gr.Checkbox(label="Keep Original Images", elem_id=self.elem_id("show_original"), visible=False, default=False)
|
|
|
|
|
|
def show_recolor_strength(rec_checked):
|
|
|
return [gr.update(visible=rec_checked), gr.update(visible=rec_checked)]
|
|
|
|
|
|
recolor.change(
|
|
|
fn=show_recolor_strength,
|
|
|
inputs=recolor,
|
|
|
outputs=[rec_strength, show_original]
|
|
|
)
|
|
|
|
|
|
self.infotext_fields = [
|
|
|
(rescale, "CFG Rescale"),
|
|
|
(recolor, "Auto Color Fix")
|
|
|
]
|
|
|
self.paste_field_names = []
|
|
|
for _, field_name in self.infotext_fields:
|
|
|
self.paste_field_names.append(field_name)
|
|
|
return [rescale, recolor, rec_strength, show_original]
|
|
|
|
|
|
def cfg_replace(self, x_out, conds_list, uncond, cond_scale):
|
|
|
denoised_uncond = x_out[-uncond.shape[0]:]
|
|
|
denoised = torch.clone(denoised_uncond)
|
|
|
fi = globals()['cfg_rescale_fi']
|
|
|
|
|
|
for i, conds in enumerate(conds_list):
|
|
|
for cond_index, weight in conds:
|
|
|
if fi == 0:
|
|
|
denoised[i] += (x_out[cond_index] - denoised_uncond[i]) * (weight * cond_scale)
|
|
|
else:
|
|
|
xcfg = (denoised_uncond[i] + (x_out[cond_index] - denoised_uncond[i]) * (cond_scale * weight))
|
|
|
xrescaled = (torch.std(x_out[cond_index]) / torch.std(xcfg))
|
|
|
xfinal = fi * xrescaled + (1.0 - fi)
|
|
|
denoised[i] = xfinal * xcfg
|
|
|
|
|
|
return denoised
|
|
|
|
|
|
def process(self, p, rescale, recolor, rec_strength, show_original):
|
|
|
|
|
|
if globals()['enable_furry_cocks']:
|
|
|
|
|
|
rescale_override = None
|
|
|
def found(m):
|
|
|
nonlocal rescale_override
|
|
|
try:
|
|
|
rescale_override = float(m.group(1))
|
|
|
except ValueError:
|
|
|
rescale_override = None
|
|
|
return ""
|
|
|
p.prompt = re.sub(re_prompt_cfgr, found, p.prompt)
|
|
|
if rescale_override is not None:
|
|
|
rescale = rescale_override
|
|
|
|
|
|
globals()['cfg_rescale_fi'] = rescale
|
|
|
else:
|
|
|
|
|
|
rescale = globals()['cfg_rescale_fi']
|
|
|
globals()['enable_furry_cocks'] = True
|
|
|
|
|
|
sd_samplers_kdiffusion.CFGDenoiser.combine_denoised = self.cfg_replace
|
|
|
|
|
|
if rescale > 0:
|
|
|
p.extra_generation_params["CFG Rescale"] = rescale
|
|
|
|
|
|
if recolor:
|
|
|
p.extra_generation_params["Auto Color Fix Strength"] = rec_strength
|
|
|
p.do_not_save_samples = True
|
|
|
|
|
|
def postprocess_batch_list(self, p, pp, rescale, recolor, rec_strength, show_original, batch_number):
|
|
|
if recolor and show_original:
|
|
|
num = len(pp.images)
|
|
|
for i in range(num):
|
|
|
pp.images.append(pp.images[i])
|
|
|
p.prompts.append(p.prompts[i])
|
|
|
p.negative_prompts.append(p.negative_prompts[i])
|
|
|
p.seeds.append(p.seeds[i])
|
|
|
p.subseeds.append(p.subseeds[i])
|
|
|
|
|
|
def postprocess(self, p, processed, rescale, recolor, rec_strength, show_original):
|
|
|
sd_samplers_kdiffusion.CFGDenoiser.combine_denoised = self.old_denoising
|
|
|
|
|
|
def postfix(img, rec_strength):
|
|
|
prec = 0.0005 * rec_strength
|
|
|
r, g, b = img.split()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
rh, rbins = np.histogram(r, 256, (0, 256))
|
|
|
tmp = np.where(rh > rh.sum() * prec)[0]
|
|
|
r_min = tmp.min()
|
|
|
r_max = tmp.max()
|
|
|
|
|
|
gh, gbins = np.histogram(g, 256, (0, 256))
|
|
|
tmp = np.where(gh > gh.sum() * prec)[0]
|
|
|
g_min = tmp.min()
|
|
|
g_max = tmp.max()
|
|
|
|
|
|
bh, bbins = np.histogram(b, 256, (0, 256))
|
|
|
tmp = np.where(bh > bh.sum() * prec)[0]
|
|
|
b_min = tmp.min()
|
|
|
b_max = tmp.max()
|
|
|
|
|
|
r = r.point(lambda i: int(255 * (min(max(i, r_min), r_max) - r_min) / (r_max - r_min)))
|
|
|
g = g.point(lambda i: int(255 * (min(max(i, g_min), g_max) - g_min) / (g_max - g_min)))
|
|
|
b = b.point(lambda i: int(255 * (min(max(i, b_min), b_max) - b_min) / (b_max - b_min)))
|
|
|
|
|
|
new_img = Image.merge("RGB", (r, g, b))
|
|
|
|
|
|
return new_img
|
|
|
|
|
|
if recolor:
|
|
|
grab = 0
|
|
|
n_img = len(processed.images)
|
|
|
for i in range(n_img):
|
|
|
doit = False
|
|
|
|
|
|
if show_original:
|
|
|
check = i
|
|
|
if opts.return_grid:
|
|
|
if i == 0:
|
|
|
continue
|
|
|
else:
|
|
|
check = check - 1
|
|
|
doit = check % (p.batch_size * 2) >= p.batch_size
|
|
|
else:
|
|
|
if n_img > 1 and i != 0:
|
|
|
doit = True
|
|
|
elif n_img == 1 or not opts.return_grid:
|
|
|
doit = True
|
|
|
|
|
|
if doit:
|
|
|
res_img = postfix(processed.images[i], rec_strength)
|
|
|
if opts.samples_save:
|
|
|
ind = grab
|
|
|
grab += 1
|
|
|
prompt_infotext = processing.create_infotext(p, p.all_prompts, p.all_seeds, p.all_subseeds,
|
|
|
index=ind)
|
|
|
|
|
|
if opts.samples_save:
|
|
|
saving.save_image(processed.images[i], p.outpath_samples, "", seed=p.all_seeds[ind],
|
|
|
prompt=p.all_prompts[ind],
|
|
|
info=prompt_infotext, p=p, suffix="colorfix")
|
|
|
saving.save_image(res_img, p.outpath_samples, "", seed=p.all_seeds[ind],
|
|
|
prompt=p.all_prompts[ind],
|
|
|
info=prompt_infotext, p=p, suffix="colorfix")
|
|
|
|
|
|
processed.images[i] = res_img
|
|
|
|
|
|
|
|
|
def on_infotext_pasted(infotext, params):
|
|
|
if "CFG Rescale" not in params:
|
|
|
params["CFG Rescale"] = 0
|
|
|
|
|
|
if "CFG Rescale φ" in params:
|
|
|
params["CFG Rescale"] = params["CFG Rescale φ"]
|
|
|
del params["CFG Rescale φ"]
|
|
|
|
|
|
if "CFG Rescale phi" in params and scripts.scripts_txt2img.script("Neutral Prompt") is None:
|
|
|
params["CFG Rescale"] = params["CFG Rescale phi"]
|
|
|
del params["CFG Rescale phi"]
|
|
|
|
|
|
if "DDIM Trailing" not in params:
|
|
|
params["DDIM Trailing"] = False
|
|
|
|
|
|
|
|
|
script_callbacks.on_infotext_pasted(on_infotext_pasted)
|
|
|
|
