From 176a269a623a2a43efefdacdf425ba5969e51eb9 Mon Sep 17 00:00:00 2001 From: Extraltodeus Date: Sun, 7 Jul 2024 15:12:48 +0200 Subject: [PATCH] Add files via upload --- nodes.py | 139 +++++++++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 119 insertions(+), 20 deletions(-) diff --git a/nodes.py b/nodes.py index 90171a5..36ae281 100644 --- a/nodes.py +++ b/nodes.py @@ -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)