From 83ad77968d164afb6c4401dbde1a964f42526a8c Mon Sep 17 00:00:00 2001 From: Extraltodeus Date: Wed, 5 Jun 2024 20:28:29 +0200 Subject: [PATCH] Add files via upload --- __init__.py | 7 +-- nodes.py | 156 +++++++++++----------------------------------------- 2 files changed, 34 insertions(+), 129 deletions(-) diff --git a/__init__.py b/__init__.py index e555ded..872b62a 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,6 @@ from .nodes import * NODE_CLASS_MAPPINGS = { - "Unet Temperature SDXL":UnetTemperaturePatchSDXL, - "Unet Temperature SD1":UnetTemperaturePatchSD15, - "Unet Temperature any model":UnetTemperaturePatchAny, - "Unet Temperature SDXL per layer":UnetTemperaturePatchSDXLpl, - "Unet Temperature SD1 per layer" :UnetTemperaturePatchSD15pl, + "Unet Temperature":UnetTemperaturePatch, "CLIP Temperature":CLIPTemperaturePatch, - "Dynamic scale as float":temperatureForScaleAsFloat, } diff --git a/nodes.py b/nodes.py index 3985d42..9d6f951 100644 --- a/nodes.py +++ b/nodes.py @@ -7,49 +7,19 @@ EPSILON = 1e-4 SD_layer_dims = { "SD1" : {"input_1": 4096,"input_2": 4096,"input_4": 1024,"input_5": 1024,"input_7": 256,"input_8": 256,"middle_0": 64,"output_3": 256,"output_4": 256,"output_5": 256,"output_6": 1024,"output_7": 1024,"output_8": 1024,"output_9": 4096,"output_10": 4096,"output_11": 4096}, "SDXL": {"input_4": 4096,"input_5": 4096,"input_7": 1024,"input_8": 1024,"middle_0": 1024,"output_0": 1024,"output_1": 1024,"output_2": 1024,"output_3": 4096,"output_4": 4096,"output_5": 4096}, + "Disabled":{} } -layers_SD15 = { - "input":[1,2,4,5,7,8], - "middle":[0], - "output":[3,4,5,6,7,8,9,10,11], -} - -layers_SDXL = { - "input":[4,5,7,8], - "middle":[0], - "output":[0,1,2,3,4,5], -} - -revert_dim = lambda x: 8 * math.sqrt(x) -printed_var = "" -def cprint(var): - global printed_var - str_var = str(var) - if printed_var != str_var: - print(" ",str_var) - printed_var = str_var - -def dynamic_scale_attention(layer_name, model_name, q_size_1, q_size_2, **kwargs): - return 1 / (math.sqrt(q_size_1) * (SD_layer_dims[model_name][layer_name] ** 0.5 / q_size_2 ** 0.5) ** 0.5) - -def dynamic_scale_attention_light(layer_name, model_name, q_size_1, q_size_2, **kwargs): - return 1 / (math.sqrt(q_size_1 * (SD_layer_dims[model_name][layer_name] ** 0.5 / q_size_2 ** 0.5) ** 0.5)) - -def temp_non_zero_div(layer_name, model_name, q_size_1, q_size_2, **kwargs): - return 1 / (math.sqrt(q_size_1) * EPSILON) - -auto_temp_methods = {"normal": dynamic_scale_attention, "light": dynamic_scale_attention_light, "clip":temp_non_zero_div} +def should_scale(mname,lname,q2): + if mname != "Disabled" and lname in SD_layer_dims[mname]: + return q2 != SD_layer_dims[mname][lname] + return False class temperature_patcher(): - def __init__(self, temperature, layer_name = "", model_name="", eval_string = "", method="medium", base_resolution=(512,512), target_resolution=(512,512)): - self.temperature = temperature + def __init__(self, temperature, layer_name = "", model_name=""): + self.temperature = max(temperature,EPSILON) self.layer_name = layer_name self.model_name = model_name - self.eval_string = eval_string - self.method = auto_temp_methods[method] - self.base_resolution = base_resolution - self.target_resolution = target_resolution 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'] @@ -60,45 +30,29 @@ class temperature_patcher(): (q, k, v), ) - if self.eval_string != "": - if self.layer_name != "": - layer_dim = SD_layer_dims[self.model_name][self.layer_name] - q_size_1 = q.size(-1) - q_size_2 = q.size(-2) - c = [] - evals_strings = self.eval_string.split(";") - if len(evals_strings) > 1: - for i in range(len(evals_strings[:-1])): - c.append(eval(evals_strings[i])) - scale = eval(evals_strings[-1]) - else: - scale = 1 / (math.sqrt(q.size(-1)) * self.temperature) if self.temperature > 0 else \ - self.method(layer_name=self.layer_name, model_name=self.model_name, q_size_1=q.size(-1), q_size_2=q.size(-2), - base_resolution=self.base_resolution,target_resolution=self.target_resolution) + scale = 1 / (math.sqrt(q.size(-1)) * self.temperature) + + 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)): + out *= math.log(q.size(-2) ** 0.5) / math.log(SD_layer_dims[self.model_name][self.layer_name] ** 0.5) - out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False,scale=scale) out = ( out.transpose(1, 2).reshape(b, -1, heads * dim_head) ) + return out class UnetTemperaturePatch: @classmethod def INPUT_TYPES(s): - if not s.ANY_MODEL: - required_inputs = {f"{key}_{layer}": ("BOOLEAN", {"default": True}) for key, layers in s.TOGGLES.items() for layer in layers} - else: - required_inputs = {} + required_inputs = {} 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"],) - if s.MODEL_NAME in ["SDXL","SD1"]: - required_inputs["DSA_intensity"] = (["normal","light"],) - # required_inputs["eval_string"] = ("STRING", {"multiline": True}) + required_inputs["Temperature"] = ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01, "round": 0.01}) + required_inputs["Attention"] = (["self","cross","both"],) + required_inputs["Dynamic_Scale_Attention"] = (["Disabled","SDXL","SD1"],) return {"required": required_inputs} - - ANY_MODEL = False - LAYER_NAME = None + TOGGLES = {} RETURN_TYPES = ("MODEL","STRING",) RETURN_NAMES = ("Model","String",) @@ -106,33 +60,23 @@ class UnetTemperaturePatch: CATEGORY = "model_patches/Temperature" - def patch(self, model, Temperature, Attention, DSA_intensity="normal", eval_string="", **kwargs): - model_name = self.__class__.MODEL_NAME - any_model = self.__class__.ANY_MODEL - - if not any_model: - layer_names = kwargs - else: - layer_names = {f"{l}_{n}": True for l in ["input", "middle", "output"] for n in range(12)} - + def patch(self, model, Temperature, Attention, Dynamic_Scale_Attention, **kwargs): m = model.clone() levels = ["input","middle","output"] - parameters_output = {level:[] for level in levels} - - for key, toggle_enabled in layer_names.items(): + layer_names = {f"{l}_{n}": True for l in levels for n in range(12)} + + for key, toggle in layer_names.items(): current_level = key.split("_")[0] - if current_level in levels and toggle_enabled: - b_number = int(key.split("_")[1]) - parameters_output[current_level].append(b_number) - patcher = temperature_patcher(Temperature,method=DSA_intensity if model_name in ["SDXL","SD1"] else "clip",layer_name=key,model_name=model_name,eval_string=eval_string) + b_number = int(key.split("_")[1]) - if Attention in ["both","self"]: - m.set_model_attn1_replace(patcher.pytorch_attention_with_temperature, current_level, b_number) - if Attention in ["both","cross"]: - m.set_model_attn2_replace(patcher.pytorch_attention_with_temperature, current_level, b_number) + if Attention in ["both","self"]: + patcher = temperature_patcher(Temperature,layer_name=key,model_name=Dynamic_Scale_Attention) + m.set_model_attn1_replace(patcher.pytorch_attention_with_temperature, current_level, b_number) + if Attention in ["both","cross"]: + patcher = temperature_patcher(Temperature,layer_name=key,model_name=Dynamic_Scale_Attention) + m.set_model_attn2_replace(patcher.pytorch_attention_with_temperature, current_level, b_number) - parameters_as_string = "\n".join(f"{k}: {','.join(map(str, v))}" for k, v in parameters_output.items()) - parameters_as_string = f"Temperature: {Temperature}\n{parameters_as_string}\nAttention: {Attention}" + parameters_as_string = f"Temperature: {Temperature}\nAttention: {Attention}\nDynamic scale: {Dynamic_Scale_Attention}" return (m, parameters_as_string,) class CLIPTemperaturePatch: @@ -147,9 +91,8 @@ class CLIPTemperaturePatch: CATEGORY = "model_patches/Temperature" def patch(self, clip, Temperature): - print(f"\n\n\nThe CLIP patch ignores the connection. Set at 1 to get default behavior. Or reload the model without this node.\n\n\n") def custom_optimized_attention(device, mask=None, small_input=True): - return temperature_patcher(Temperature,method="clip").pytorch_attention_with_temperature + return temperature_patcher(Temperature).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) @@ -172,37 +115,4 @@ class CLIPTemperaturePatch: clip_encoder_instance_g = clip.cond_stage_model.clip_g.transformer.text_model.encoder clip_encoder_instance_g.forward = types.MethodType(new_forward, clip_encoder_instance_g) - return (clip,) - -class temperatureForScaleAsFloat: - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "base_resolution" : ("INT", {"default": 1024, "min": 0,"max": 10000,"step": 8}), - "target_width" : ("INT", {"default": 1024, "min": 0,"max": 10000,"step": 8}), - "target_height" : ("INT", {"default": 1024, "min": 0,"max": 10000,"step": 8}), - "multiply_value_by" : ("FLOAT", {"default": 1, "min": 0, "max": 100,"step": 0.01}), - "shift_value_by" : ("FLOAT", {"default": 0, "min": -100,"max": 100,"step": 0.01}), - "print_value" : ("BOOLEAN", {"default": False}), - } - } - - FUNCTION = "simple_output" - RETURN_TYPES = ("FLOAT",) - CATEGORY = "model_patches/Temperature" - - def simple_output(self, base_resolution,target_width,target_height,multiply_value_by,shift_value_by,print_value): - dsa = ((base_resolution / (target_width*target_height) ** 0.5) ** 0.5) * multiply_value_by + shift_value_by - if print_value: - print(f"\nDynamic scale attention is {dsa}\n") - return (dsa,) - -UnetTemperaturePatchSDXL = type("Unet Temperature SDXL", (UnetTemperaturePatch,), {"TOGGLES": layers_SDXL,"MODEL_NAME":"SDXL","ANY_MODEL": True}) -UnetTemperaturePatchSD15 = type("Unet Temperature SD1", (UnetTemperaturePatch,), {"TOGGLES": layers_SD15,"MODEL_NAME":"SD1", "ANY_MODEL": True,}) -UnetTemperaturePatchSDXLpl = type("Unet Temperature SDXL per layer", (UnetTemperaturePatch,), {"TOGGLES": layers_SDXL,"MODEL_NAME":"SDXL","ANY_MODEL": False}) -UnetTemperaturePatchSD15pl = type("Unet Temperature SD1 per layer", (UnetTemperaturePatch,), {"TOGGLES": layers_SD15,"MODEL_NAME":"SD1", "ANY_MODEL": False}) -UnetTemperaturePatchAny = type("Unet Temperature any model", (UnetTemperaturePatch,), {"MODEL_NAME":"","ANY_MODEL": True}) + return (clip,) \ No newline at end of file