diff --git a/TGate.py b/TGate.py new file mode 100644 index 0000000..78ab98f --- /dev/null +++ b/TGate.py @@ -0,0 +1,194 @@ +import torch + + +class TGateTransformerWrapper: + def __init__(self, sigma_gate_attn1=-1, sigma_gate_attn2=-1): + self.sigma_gate_attn1 = sigma_gate_attn1 + self.sigma_gate_attn2 = sigma_gate_attn2 + self.attn1_cache = {} + self.attn2_cache = {} + + def __call__(self, inner_block, x, context=None, transformer_options=None): + transformer_options = transformer_options or {} + extra_options = {} + block = transformer_options.get("block", None) + block_index = transformer_options.get("block_index", 0) + transformer_patches = {} + transformer_patches_replace = {} + + for k in transformer_options: + if k == "patches": + transformer_patches = transformer_options[k] + elif k == "patches_replace": + transformer_patches_replace = transformer_options[k] + else: + extra_options[k] = transformer_options[k] + + extra_options["n_heads"] = inner_block.n_heads + extra_options["dim_head"] = inner_block.d_head + sigma = extra_options["sigmas"].detach().cpu()[0].item() if "sigmas" in extra_options else 999999999.9 + # transformer_block key + transformer_block = (block[0], block[1], block_index) if block is not None else None + + # in resblock + if inner_block.ff_in: + x_skip = x + x = inner_block.ff_in(inner_block.norm_in(x)) + if inner_block.is_res: + x += x_skip + + n = inner_block.norm1(x) + if sigma < self.sigma_gate_attn1 and transformer_block in self.attn1_cache: + # use cache + n = self.attn1_cache[transformer_block] + if sigma < self.sigma_gate_attn2: + uncond, cond = n.chunk(2) + n = (uncond + cond) / 2 + # n = torch.mean(n, dim=0, keepdim=True) + else: + context_attn1 = context if inner_block.disable_self_attn else None + value_attn1 = None + + if "attn1_patch" in transformer_patches: + patch = transformer_patches["attn1_patch"] + if context_attn1 is None: + context_attn1 = n + value_attn1 = context_attn1 + for p in patch: + n, context_attn1, value_attn1 = p(n, context_attn1, value_attn1, extra_options) + + attn1_replace_patch = transformer_patches_replace.get("attn1", {}) + block_attn1 = transformer_block + if block_attn1 not in attn1_replace_patch: + block_attn1 = block + + # do attn + if block_attn1 in attn1_replace_patch: + if context_attn1 is None: + context_attn1 = n + value_attn1 = n + n = inner_block.attn1.to_q(n) + context_attn1 = inner_block.attn1.to_k(context_attn1) + value_attn1 = inner_block.attn1.to_v(value_attn1) + n = attn1_replace_patch[block_attn1](n, context_attn1, value_attn1, extra_options) + n = inner_block.attn1.to_out(n) + else: + n = inner_block.attn1(n, context=context_attn1, value=value_attn1) + + if "attn1_output_patch" in transformer_patches: + patch = transformer_patches["attn1_output_patch"] + for p in patch: + n = p(n, extra_options) + + self.attn1_cache[transformer_block] = n # TODO: reduce, cache times + + x += n + if "middle_patch" in transformer_patches: + patch = transformer_patches["middle_patch"] + for p in patch: + x = p(x, extra_options) + + if inner_block.attn2 is not None: + n = inner_block.norm2(x) + context_attn2 = n if inner_block.switch_temporal_ca_to_sa else context + value_attn2 = None + if "attn2_patch" in transformer_patches: + patch = transformer_patches["attn2_patch"] + value_attn2 = context_attn2 + for p in patch: + n, context_attn2, value_attn2 = p(n, context_attn2, value_attn2, extra_options) + + if sigma < self.sigma_gate_attn2 and transformer_block in self.attn2_cache: + n = self.attn2_cache[transformer_block] + if sigma < self.sigma_gate_attn1: + uncond, cond = n.chunk(2) + n = (uncond + cond) / 2 + # n = torch.mean(n, dim=0, keepdim=True) + else: + attn2_replace_patch = transformer_patches_replace.get("attn2", {}) + block_attn2 = transformer_block + if block_attn2 not in attn2_replace_patch: + block_attn2 = block + + if block_attn2 in attn2_replace_patch: + if value_attn2 is None: + value_attn2 = context_attn2 + n = inner_block.attn2.to_q(n) + context_attn2 = inner_block.attn2.to_k(context_attn2) + value_attn2 = inner_block.attn2.to_v(value_attn2) + n = attn2_replace_patch[block_attn2](n, context_attn2, value_attn2, extra_options) + n = inner_block.attn2.to_out(n) + else: + n = inner_block.attn2(n, context=context_attn2, value=value_attn2) + if "attn2_output_patch" in transformer_patches: + patch = transformer_patches["attn2_output_patch"] + for p in patch: + n = p(n, extra_options) + self.attn2_cache[transformer_block] = n + + x += n + if inner_block.is_res: + x_skip = x + x = inner_block.ff(inner_block.norm3(x)) + if inner_block.is_res: + x += x_skip + + return x + + +class TGateSamplerCfgRescaler: + def __init__(self, sigma_gate_attn1=-1, sigma_gate_attn2=-1): + self.sigma_gate = min(sigma_gate_attn1, sigma_gate_attn2) + + def __call__(self, kwds): + sigma = kwds["timestep"].detach().cpu()[0].item() if "timestep" in kwds else 999999999.9 + if sigma < self.sigma_gate: + return 1.0 + return kwds["cond_scale"] + + +class TGateSamplerFn: + def __init__(self, sigma_gate_attn1=-1, sigma_gate_attn2=-1): + self.sigma_gate = min(sigma_gate_attn1, sigma_gate_attn2) + + def __call__(self, kwds): + sigma = kwds["timestep"].detach().cpu()[0].item() if "timestep" in kwds else 999999999.9 + if sigma < self.sigma_gate: + return kwds["cond"] * kwds["cond_scale"] + # uncond_pred + (cond_pred - uncond_pred) * cond_scale + # x - [x - uncond_pred + (x - uncond_pred - (x - cond_pred)) * cond_scale] + return kwds["uncond"] + (kwds["uncond"] - kwds["cond"]) * kwds["cond_scale"] + + +class TGateApply: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("MODEL",), + "attn1_start_at": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "attn2_start_at": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + }, + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "apply_tgate" + CATEGORY = "TGate" + + def apply_tgate(self, model, attn1_start_at=1.0, attn2_start_at=1.0): + model_clone = model.clone() + sigma_attn1 = model_clone.get_model_object("model_sampling").percent_to_sigma(attn1_start_at) + sigma_attn2 = model_clone.get_model_object("model_sampling").percent_to_sigma(attn2_start_at) + model_clone.set_model_transformer_function(TGateTransformerWrapper(sigma_attn1, sigma_attn2)) + model_clone.set_model_sampler_cfg_rescaler(TGateSamplerCfgRescaler(sigma_attn1, sigma_attn2)) + # model_clone.set_model_sampler_cfg_function(TGateSamplerFn(sigma_attn1, sigma_attn2)) + return (model_clone,) + + +NODE_CLASS_MAPPINGS = { + "TGateApply": TGateApply, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "TGateApply": "TGate Apply", +} diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..a2d69ca --- /dev/null +++ b/__init__.py @@ -0,0 +1 @@ +from .TGate import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS diff --git a/tgate.patch b/tgate.patch new file mode 100644 index 0000000..0cbf6ca --- /dev/null +++ b/tgate.patch @@ -0,0 +1,45 @@ +diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py +index f116efee..e52466b3 100644 +--- a/comfy/ldm/modules/attention.py ++++ b/comfy/ldm/modules/attention.py +@@ -460,6 +460,8 @@ class BasicTransformerBlock(nn.Module): + return checkpoint(self._forward, (x, context, transformer_options), self.parameters(), self.checkpoint) + + def _forward(self, x, context=None, transformer_options={}): ++ if "transformer_function" in transformer_options: ++ return transformer_options["transformer_function"](self, x, context, transformer_options) + extra_options = {} + block = transformer_options.get("block", None) + block_index = transformer_options.get("block_index", 0) +diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py +index cf51c4ad..21d5b1af 100644 +--- a/comfy/model_patcher.py ++++ b/comfy/model_patcher.py +@@ -158,6 +158,14 @@ class ModelPatcher: + def set_model_output_block_patch(self, patch): + self.set_model_patch(patch, "output_block_patch") + ++ def set_model_transformer_function(self, transformer_function): ++ self.model_options["transformer_options"]["transformer_function"] = transformer_function ++ ++ def set_model_sampler_cfg_rescaler(self, cfg_rescaler_fn, disable_cfg1_optimization=False): ++ self.model_options["sampler_cfg_rescaler"] = cfg_rescaler_fn ++ if disable_cfg1_optimization: ++ self.model_options["disable_cfg1_optimization"] = True ++ + def add_object_patch(self, name, obj): + self.object_patches[name] = obj + +diff --git a/comfy/samplers.py b/comfy/samplers.py +index 415a35cc..bd8912bc 100644 +--- a/comfy/samplers.py ++++ b/comfy/samplers.py +@@ -249,6 +249,8 @@ def cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_o + #The main sampling function shared by all the samplers + #Returns denoised + def sampling_function(model, x, timestep, uncond, cond, cond_scale, model_options={}, seed=None): ++ if "sampler_cfg_rescaler" in model_options: ++ cond_scale = model_options["sampler_cfg_rescaler"]({"cond_scale": cond_scale, "timestep": timestep}) + if math.isclose(cond_scale, 1.0) and model_options.get("disable_cfg1_optimization", False) == False: + uncond_ = None + else: