nodes/Extraltodeus-Stable-Diffusion-temperature-settings
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
Extraltodeus
2024-07-07 15:12:48 +02:00
committed by GitHub
parent 1ec74c54bd
commit 176a269a62
1 changed files with 119 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+119 -20
View File
@@ -1,8 +1,9 @@
from functools import partial
import torch
import math
import types
import comfy.model_management
from functools import partial
from math import *
from comfy.model_patcher import ModelPatcher
from comfy.sd import CLIP
@@ -16,18 +17,45 @@ SD_layer_dims = {
models_by_size = {"1719049928": "SD1", "5134967368":"SDXL"}
def cv_temperature(input_tensor, auto_mode="normal"):
if "creative" in auto_mode:
tensor_std = input_tensor.std()
temperature = (torch.std(torch.abs(input_tensor - tensor_std))/tensor_std)
temperature = 1 / temperature
del tensor_std
else:
temperature = torch.std(input_tensor)
if "squared" in auto_mode:
temperature = temperature ** 2
elif "sqrt" in auto_mode:
temperature = temperature ** .5
if not "reversed" in auto_mode:
temperature = 1 / temperature
return temperature
def should_scale(mname,lname,q2):
if mname == "": return False
if mname == "" or mname == "CLIP": return False
if mname != "Disabled" and lname in SD_layer_dims[mname]:
if lname not in SD_layer_dims[mname]:
return False
return q2 != SD_layer_dims[mname][lname]
return False
class temperature_patcher():
def __init__(self, temperature, layer_name = "", model_name="", eval_string=""):
def __init__(self, temperature, layer_name = "", model_name="", eval_string="", auto_temp="disabled",
Original_scale=512, Target_scale_X=512, Target_scale_Y=512, rescale_adjust=1,
scale_before=False,scale_after=False):
self.temperature = max(temperature,EPSILON)
self.layer_name = layer_name
self.model_name = model_name
self.eval_string = eval_string
self.auto_temp = auto_temp
self.Original_scale = Original_scale
self.Target_scale_X = Target_scale_X
self.Target_scale_Y = Target_scale_Y
self.rescale_adjust = rescale_adjust
self.scale_before = scale_before
self.scale_after = scale_after
def pytorch_attention_with_temperature(self, q, k, v, extra_options, mask=None, attn_precision=None):
heads = extra_options if isinstance(extra_options, int) else extra_options['n_heads']
@@ -38,20 +66,42 @@ class temperature_patcher():
(q, k, v),
)
scale = 1 / (math.sqrt(q.size(-1)) * self.temperature)
if self.auto_temp != "disabled":
extra_temperature = cv_temperature({'q_': q, 'k_': k, 'v_': v}[self.auto_temp[:2]], self.auto_temp)
else:
extra_temperature = 1
temperature_pre_scale = 1
if self.scale_before:
if should_scale(self.model_name, self.layer_name,q.size(-2)):
ldim = SD_layer_dims[self.model_name][self.layer_name]
if self.eval_string != "":
temperature_pre_scale = eval(self.eval_string)
else:
temperature_pre_scale = log(q.size(-2), ldim) # that's the actual resccale, everything else is for testing purpose
elif (self.Target_scale_X*self.Target_scale_Y) != self.Original_scale**2:
temperature_pre_scale = log((self.Target_scale_X*self.Target_scale_Y)**.5, self.Original_scale)
temperature_scale = self.temperature / temperature_pre_scale
out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False, scale=scale)
if should_scale(self.model_name, self.layer_name,q.size(-2)):
ldim = SD_layer_dims[self.model_name][self.layer_name]
if self.eval_string != "":
out = eval(self.eval_string)
else:
out *= math.log(q.size(-2), ldim)
scale = 1 / (math.sqrt(q.size(-1)) * temperature_scale * extra_temperature)
out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, is_causal=False, scale=scale)
if self.scale_after:
if should_scale(self.model_name, self.layer_name,q.size(-2)):
ldim = SD_layer_dims[self.model_name][self.layer_name]
if self.eval_string != "":
out *= eval(self.eval_string)
else:
out *= log(q.size(-2), ldim)
elif (self.Target_scale_X*self.Target_scale_Y) != self.Original_scale**2:
out *= log((self.Target_scale_X*self.Target_scale_Y)**.5, self.Original_scale)
if self.rescale_adjust != 1:
out *= self.rescale_adjust
out = (
out.transpose(1, 2).reshape(b, -1, heads * dim_head)
)
return out
class UnetTemperaturePatch:
@@ -61,7 +111,8 @@ class UnetTemperaturePatch:
required_inputs["model"] = ("MODEL",)
required_inputs["Temperature"] = ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01, "round": 0.01})
required_inputs["Attention"] = (["both","self","cross"],)
required_inputs["Dynamic_Scale_Attention"] = ("BOOLEAN", {"default": True})
required_inputs["Dynamic_Scale_Temperature"] = ("BOOLEAN", {"default": False})
required_inputs["Dynamic_Scale_Output"] = ("BOOLEAN", {"default": True})
# required_inputs["eval_string"] = ("STRING", {"multiline": True})
return {"required": required_inputs}
@@ -72,7 +123,11 @@ class UnetTemperaturePatch:
CATEGORY = "model_patches/Temperature"
def patch(self, model: ModelPatcher, Temperature, Attention, Dynamic_Scale_Attention, eval_string="", **kwargs):
def patch(self, model, Temperature, Attention, Dynamic_Scale_Temperature, Dynamic_Scale_Output, Dynamic_Scale_Adjust=1, eval_string="", **kwargs):
Dynamic_Scale_Attention = Dynamic_Scale_Temperature or Dynamic_Scale_Output
if not Dynamic_Scale_Attention and Temperature == 1:
print("Dynamic_Scale_Attention/temperature: no patch applied.")
return (model, "Fully disabled",)
if Dynamic_Scale_Attention and str(model.size) in models_by_size:
model_name = models_by_size[str(model.size)]
print(f"Model detected for scaling: {model_name}")
@@ -88,7 +143,9 @@ class UnetTemperaturePatch:
for key, toggle in layer_names.items():
current_level = key.split("_")[0]
b_number = int(key.split("_")[1])
patcher = temperature_patcher(Temperature,layer_name=key,model_name=model_name, eval_string=eval_string)
patcher = temperature_patcher(Temperature,layer_name=key,model_name=model_name, eval_string=eval_string,
rescale_adjust=Dynamic_Scale_Adjust, scale_before=Dynamic_Scale_Temperature,
scale_after=Dynamic_Scale_Output)
if Attention in ["both","self"]:
m.set_model_attn1_replace(patcher.pytorch_attention_with_temperature, current_level, b_number)
if Attention in ["both","cross"]:
@@ -102,18 +159,60 @@ class CLIPTemperaturePatch:
def INPUT_TYPES(cls):
return {"required": { "clip": ("CLIP",),
"Temperature": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
# "Auto_temp": ("BOOLEAN", {"default": False}) # It's just experimental but uncomment it if you want to try.
}}
RETURN_TYPES = ("CLIP",)
FUNCTION = "patch"
CATEGORY = "model_patches/Temperature"
def patch(self, clip: CLIP, Temperature):
def patch(self, clip, Temperature, Auto_temp=False):
c = clip.clone()
def custom_optimized_attention(device, mask=None, small_input=True):
return temperature_patcher(Temperature).pytorch_attention_with_temperature
return temperature_patcher(temperature=Temperature,auto_temp="k_creative" if Auto_temp else "disabled").pytorch_attention_with_temperature
def new_forward(self, x, mask=None, intermediate_output=None):
optimized_attention = custom_optimized_attention(x.device, mask=mask is not None, small_input=True)
if intermediate_output is not None:
if intermediate_output < 0:
intermediate_output = len(self.layers) + intermediate_output
intermediate = None
for i, l in enumerate(self.layers):
x = l(x, mask, optimized_attention)
if i == intermediate_output:
intermediate = x.clone()
return x, intermediate
if getattr(c.patcher.model, "clip_g", None) is not None:
c.patcher.add_object_patch("clip_g.transformer.text_model.encoder.forward", partial(new_forward, c.patcher.model.clip_g.transformer.text_model.encoder))
if getattr(c.patcher.model, "clip_l", None) is not None:
c.patcher.add_object_patch("clip_l.transformer.text_model.encoder.forward", partial(new_forward, c.patcher.model.clip_l.transformer.text_model.encoder))
return (c,)
class CLIPTemperatureWithScalePatch:
@classmethod
def INPUT_TYPES(cls):
return {"required": { "clip": ("CLIP",),
"Temperature": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
"Original_scale": ("INT", {"default": 512, "min": 64, "max": 100000.0, "step": 64}),
"Target_scale_X": ("INT", {"default": 512, "min": 64, "max": 100000.0, "step": 64}),
"Target_scale_Y": ("INT", {"default": 512, "min": 64, "max": 100000.0, "step": 64}),
"Dynamic_Scale_Temperature": ("BOOLEAN", {"default": False}),
"Dynamic_Scale_Output": ("BOOLEAN", {"default": False})
}}
RETURN_TYPES = ("CLIP",)
FUNCTION = "patch"
CATEGORY = "model_patches/Temperature"
def patch(self, clip, Temperature, Dynamic_Scale_Temperature, Dynamic_Scale_Output, Original_scale=512, Target_scale_X=512, Target_scale_Y=512, Scale_Adjust=1):
c = clip.clone()
def custom_optimized_attention(device, mask=None, small_input=True):
return temperature_patcher(temperature=Temperature,Target_scale_X=Target_scale_X,Target_scale_Y=Target_scale_Y,Original_scale=Original_scale, rescale_adjust=Scale_Adjust, scale_before=Dynamic_Scale_Temperature, scale_after=Dynamic_Scale_Output).pytorch_attention_with_temperature
def new_forward(self, x, mask=None, intermediate_output=None):
optimized_attention = custom_optimized_attention(x.device, mask=mask is not None, small_input=True)