mirror of
https://github.com/lllyasviel/stable-diffusion-webui-forge.git
synced 2026-07-21 21:01:24 +08:00
63 lines
2.0 KiB
Python
63 lines
2.0 KiB
Python
# https://github.com/comfyanonymous/ComfyUI/blob/v0.3.64/comfy_extras/nodes_mahiro.py
|
|
|
|
import gradio as gr
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from modules import scripts
|
|
from modules.infotext_utils import PasteField
|
|
from modules.shared import opts
|
|
|
|
|
|
class ScriptMahiro(scripts.ScriptBuiltinUI):
|
|
section = "cfg"
|
|
create_group = False
|
|
sorting_priority = 1
|
|
|
|
def title(self):
|
|
return "MaHiRo"
|
|
|
|
def show(self, is_img2img):
|
|
return scripts.AlwaysVisible if opts.show_mahiro else None
|
|
|
|
def ui(self, is_img2img):
|
|
enable = gr.Checkbox(
|
|
value=False,
|
|
label="MaHiRo",
|
|
elem_id=f"{'img2img' if is_img2img else 'txt2img'}_enable_mahiro",
|
|
scale=1,
|
|
)
|
|
self.infotext_fields = [PasteField(enable, "MaHiRo", api="mahiro")]
|
|
return [enable]
|
|
|
|
def after_extra_networks_activate(self, p, enable, *args, **kwargs):
|
|
if enable:
|
|
p.extra_generation_params.update({"MaHiRo": enable})
|
|
p.mahiro = True
|
|
|
|
def process_before_every_sampling(self, p, *args, **kwargs):
|
|
if str(getattr(p, "mahiro", False)) != "True":
|
|
return
|
|
|
|
@torch.inference_mode()
|
|
def mahiro_normd(args: dict):
|
|
scale: float = args["cond_scale"]
|
|
cond_p: torch.Tensor = args["cond_denoised"]
|
|
uncond_p: torch.Tensor = args["uncond_denoised"]
|
|
leap = cond_p * scale
|
|
u_leap = uncond_p * scale
|
|
cfg: torch.Tensor = args["denoised"]
|
|
merge = (leap + cfg) / 2
|
|
normu = torch.sqrt(u_leap.abs()) * u_leap.sign()
|
|
normm = torch.sqrt(merge.abs()) * merge.sign()
|
|
sim = F.cosine_similarity(normu, normm).mean()
|
|
simsc = 2 * (sim + 1)
|
|
wm = (simsc * cfg + (4 - simsc) * leap) / 4
|
|
return wm
|
|
|
|
unet = p.sd_model.forge_objects.unet.clone()
|
|
unet.set_model_sampler_post_cfg_function(mahiro_normd)
|
|
p.sd_model.forge_objects.unet = unet
|
|
|
|
print("using MaHiRo")
|