From 19d7d22a882584bd5cbc086aa0c7ca7523a14248 Mon Sep 17 00:00:00 2001 From: logtd Date: Sat, 24 Aug 2024 17:20:21 -0600 Subject: [PATCH] init --- .gitignore | 3 ++ __init__.py | 19 +++++++ modules/clip_vision_override.py | 26 +++++++++ modules/ssr_modules.py | 89 +++++++++++++++++++++++++++++++ nodes/apply_ssr_attention_node.py | 27 ++++++++++ nodes/encode_ssr_query_node.py | 70 ++++++++++++++++++++++++ nodes/load_ssr_aligner_node.py | 30 +++++++++++ nodes/load_ssr_attention_node.py | 57 ++++++++++++++++++++ 8 files changed, 321 insertions(+) create mode 100644 .gitignore create mode 100644 __init__.py create mode 100644 modules/clip_vision_override.py create mode 100644 modules/ssr_modules.py create mode 100644 nodes/apply_ssr_attention_node.py create mode 100644 nodes/encode_ssr_query_node.py create mode 100644 nodes/load_ssr_aligner_node.py create mode 100644 nodes/load_ssr_attention_node.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..65a35fb --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +**/__pycache__ +**/*.pyc +.DS_Store \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..25e47c0 --- /dev/null +++ b/__init__.py @@ -0,0 +1,19 @@ +from .nodes.apply_ssr_attention_node import ApplySSRNode +from .nodes.encode_ssr_query_node import EncodeSSRQueryNode +from .nodes.load_ssr_aligner_node import LoadSSRAlignerNode +from .nodes.load_ssr_attention_node import LoadSSRAttentionNode + + +NODE_CLASS_MAPPINGS = { + 'ApplySSR': ApplySSRNode, + 'EncodeSSRQuery': EncodeSSRQueryNode, + 'LoadSSRAligner': LoadSSRAlignerNode, + 'LoadSSRAttention': LoadSSRAttentionNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + 'ApplySSR': 'SSR Apply Model', + 'EncodeSSRQuery': 'SSR Encode Query', + 'LoadSSRAligner': 'SSR Load Aligner', + 'LoadSSRAttention': 'SSR Load Attentions', +} diff --git a/modules/clip_vision_override.py b/modules/clip_vision_override.py new file mode 100644 index 0000000..621d71a --- /dev/null +++ b/modules/clip_vision_override.py @@ -0,0 +1,26 @@ +from comfy.clip_vision import clip_preprocess +from comfy.ldm.modules.attention import optimized_attention_for_device +import comfy.model_management + + +def encoder(encoder, x, mask=None): + optimized_attention = optimized_attention_for_device(x.device, mask=mask is not None, small_input=True) + intermediates = [] + for i, l in enumerate(encoder.layers): + x = l(x, mask, optimized_attention) + intermediates.append(x.clone()) + return intermediates + + +def vision_model(vision_model, pixel_values): + x = vision_model.embeddings(pixel_values) + x = vision_model.pre_layrnorm(x) + intermediates = encoder(vision_model.encoder, x) + return intermediates + + +def encode_image(clip_vision, image): + comfy.model_management.load_model_gpu(clip_vision.patcher) + pixel_values = clip_preprocess(image.to(clip_vision.load_device), size=clip_vision.image_size).float() + intermediates = vision_model(clip_vision.model.vision_model, pixel_values) + return intermediates diff --git a/modules/ssr_modules.py b/modules/ssr_modules.py new file mode 100644 index 0000000..49ef4e6 --- /dev/null +++ b/modules/ssr_modules.py @@ -0,0 +1,89 @@ +import torch +import torch.nn.functional as F +from torch import nn + +import comfy.ops +ops = comfy.ops.disable_weight_init +from comfy.ldm.modules.attention import optimized_attention + + +class SSRAligner(nn.Module): + def __init__(self) -> None: + super().__init__() + query_dim = 768 + inner_dim = 512 + cross_dim = 1024 + self.heads = 8 + + self.to_q = ops.Linear(query_dim, inner_dim, bias=False) + self.to_k = ops.Linear(cross_dim, inner_dim, bias=False) + + self.to_v = nn.ModuleList([]) + self.to_out = nn.ModuleList([]) + for _ in range(6): + self.to_v.append(ops.Linear(cross_dim, inner_dim, bias=False)) + self.to_out.append(ops.Linear(inner_dim, query_dim, bias=False)) + self.norm = ops.LayerNorm(query_dim) + + def forward(self, x, context_list): + q = self.to_q(x) + k = self.to_k(context_list[-1]) + + outs = [] + for i in range(6): + v = self.to_v[i](context_list[i]) + attn_out = optimized_attention(q, k, v, self.heads) + outs.append(self.to_out[i](attn_out)) + + return self.norm(torch.cat(outs, dim=1)) # [1, 77*6, 768] + + +# me being lazy +ssr_caa_dims = { + 'input': [320, 320, 640, 640, 1280, 1280], + 'middle': [1280], + 'output': [1280, 1280, 1280, 640, 640, 640, 320, 320, 320] +} + + +class SSRCrossAttention(nn.Module): + def __init__(self, hidden_size) -> None: + super().__init__() + self.heads = 8 + + self.to_k = ops.Linear(768, hidden_size, bias=False) + self.to_v = ops.Linear(768, hidden_size, bias=False) + + self.pos_contexts = [] # [{ 'scale': int, 'embed': tensor }, ...] + self.neg_contexts = [] + + def to(self, device): + super().to(device) + for i in range(len(self.pos_contexts)): + self.pos_contexts[i]['embed'] = self.pos_contexts[i]['embed'].to(device) + + for i in range(len(self.neg_contexts)): + self.neg_contexts[i]['embed'] = self.neg_contexts[i]['embed'].to(device) + + return self + + def forward(self, q, k, v, extra_options): + + hidden_states = optimized_attention(q, k, v, extra_options['n_heads']) + cond_or_uncond = extra_options['cond_or_uncond'] + len_conds = len(cond_or_uncond) + batch_size = len(q)//len_conds + + for cond_idx, cond in enumerate(cond_or_uncond): + contexts = self.pos_contexts if cond == 0 else self.neg_contexts + q_cond = q[cond_idx*batch_size:(cond_idx+1)*batch_size] + for context in contexts: + scale = context['scale'] + embed = context['embed'] + ssr_k = self.to_k(embed).repeat(batch_size, 1, 1) + ssr_v = self.to_v(embed).repeat(batch_size, 1, 1) + ssr_hidden_states = optimized_attention(q_cond, ssr_k, ssr_v, self.heads) + + hidden_states[cond_idx*batch_size:(cond_idx+1)*batch_size] += scale * ssr_hidden_states + + return hidden_states diff --git a/nodes/apply_ssr_attention_node.py b/nodes/apply_ssr_attention_node.py new file mode 100644 index 0000000..79a0685 --- /dev/null +++ b/nodes/apply_ssr_attention_node.py @@ -0,0 +1,27 @@ + +class ApplySSRNode: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "ssr_attentions": ("SSR_ATTENTIONS",), + "ssr_embeds": ("SSR_EMBEDS",), + } + } + RETURN_TYPES = ("MODEL",) + FUNCTION = "apply" + + CATEGORY = "ssr" + + def apply(self, model, ssr_attentions, ssr_embeds): + model = model.clone() + + for key, ssr_attention in ssr_attentions.items(): + ssr_attention.pos_contexts = ssr_embeds['positive'] + ssr_attention.neg_contexts = ssr_embeds['negative'] + block_name, block_number = key + model.set_model_attn2_replace(ssr_attention, block_name, block_number) + + return (model,) + diff --git a/nodes/encode_ssr_query_node.py b/nodes/encode_ssr_query_node.py new file mode 100644 index 0000000..927d8b9 --- /dev/null +++ b/nodes/encode_ssr_query_node.py @@ -0,0 +1,70 @@ +import comfy.model_management + +from ..modules.clip_vision_override import encode_image + + +class EncodeSSRQueryNode: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "positive_scale": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 10.0, "step": 0.01}), + "positive": ("STRING", {"multiline": True, "dynamicPrompts": True, "tooltip": "The positive image query"}), + "negative_scale": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "negative": ("STRING", {"multiline": True, "dynamicPrompts": True, "tooltip": "The negative image query"}), + "clip": ("CLIP", {"tooltip": "The CLIP model used for encoding the text."}), + "clip_vision": ("CLIP_VISION",), + "ssr_aligner": ("SSR_ALIGNER",), + "image": ("IMAGE",), + }, + "optional": { + "prev_ssr_embeds": ("SSR_EMBEDS",) + } + } + RETURN_TYPES = ("SSR_EMBEDS",) + FUNCTION = "encode" + CATEGORY = "ssr" + + def _encode_text(self, text, clip): + tokens = clip.tokenize(text) + output = clip.encode_from_tokens(tokens, return_pooled=True, return_dict=True) + cond = output.pop("cond") + return cond + + def encode(self, + positive_scale, + positive, + negative_scale, + negative, + clip, + clip_vision, + ssr_aligner, + image, + prev_ssr_embeds=None): + if prev_ssr_embeds is None: + prev_ssr_embeds = { 'negative': [], 'positive': [] } + else: + prev_ssr_embeds = { + 'negative': [*prev_ssr_embeds['negative']], + 'positive': [*prev_ssr_embeds['positive']] + } + + + intermediate_device = comfy.model_management.intermediate_device() + torch_device = comfy.model_management.get_torch_device() + image_embeds = encode_image(clip_vision, image)[3::4] + + ssr_aligner.to(torch_device) + pos_embed = None + if positive and positive_scale > 0: + pos_embed = self._encode_text(positive, clip).to(torch_device) + pos_embed = ssr_aligner(pos_embed, image_embeds).to(intermediate_device) + prev_ssr_embeds['positive'].append({ 'scale': positive_scale, 'embed': pos_embed }) + if negative and negative_scale > 0: + neg_embed = self._encode_text(negative, clip).to(torch_device) + neg_embed = ssr_aligner(neg_embed, image_embeds).to(intermediate_device) + prev_ssr_embeds['negative'].append({ 'scale': negative_scale, 'embed': neg_embed }) + + ssr_aligner.to(intermediate_device) + + return (prev_ssr_embeds, ) \ No newline at end of file diff --git a/nodes/load_ssr_aligner_node.py b/nodes/load_ssr_aligner_node.py new file mode 100644 index 0000000..132f23d --- /dev/null +++ b/nodes/load_ssr_aligner_node.py @@ -0,0 +1,30 @@ +import os + +from safetensors.torch import load_file + +import folder_paths +from folder_paths import models_dir + +from ..modules.ssr_modules import SSRAligner + + +SSR_PATH = os.path.join(models_dir, 'ssr') +os.makedirs(SSR_PATH, exist_ok=True) + + +class LoadSSRAlignerNode: + @classmethod + def INPUT_TYPES(s): + return {"required": { "checkpoint": (os.listdir(SSR_PATH), )}} + + RETURN_TYPES = ("SSR_ALIGNER",) + FUNCTION = "load" + + CATEGORY = "ssr" + + def load(self, checkpoint): + checkpoint_path = os.path.join(SSR_PATH, checkpoint) + state_dict = load_file(checkpoint_path) + ssr_aligner = SSRAligner() + ssr_aligner.load_state_dict(state_dict) + return (ssr_aligner,) \ No newline at end of file diff --git a/nodes/load_ssr_attention_node.py b/nodes/load_ssr_attention_node.py new file mode 100644 index 0000000..d4aea0b --- /dev/null +++ b/nodes/load_ssr_attention_node.py @@ -0,0 +1,57 @@ +import os + +from safetensors.torch import load_file + +from folder_paths import models_dir + +from ..modules.ssr_modules import SSRCrossAttention, ssr_caa_dims + + +SSR_PATH = os.path.join(models_dir, 'ssr') +os.makedirs(SSR_PATH, exist_ok=True) + + +class LoadSSRAttentionNode: + @classmethod + def INPUT_TYPES(s): + return {"required": { "checkpoint": (os.listdir(SSR_PATH), )}} + + RETURN_TYPES = ("SSR_ATTENTIONS",) + FUNCTION = "load" + + CATEGORY = "ssr" + + def load(self, checkpoint): + checkpoint_path = os.path.join(SSR_PATH, checkpoint) + state_dict = load_file(checkpoint_path) + ssr_idx = 1 + ssr_attentions = {} + for idx, i in enumerate([1,2,3,4,5,6]): + ssr_aligner = SSRCrossAttention(ssr_caa_dims['input'][idx]) + attn_dict = { + 'to_k.weight': state_dict[f'{ssr_idx}.to_k_SSR.weight'], + 'to_v.weight': state_dict[f'{ssr_idx}.to_v_SSR.weight'], + } + ssr_aligner.load_state_dict(attn_dict) + ssr_attentions[('input', i)] = ssr_aligner + ssr_idx += 2 + + for idx, i in enumerate([3,4,5,6,7,8,9,10,11]): + ssr_aligner = SSRCrossAttention(ssr_caa_dims['output'][idx]) + attn_dict = { + 'to_k.weight': state_dict[f'{ssr_idx}.to_k_SSR.weight'], + 'to_v.weight': state_dict[f'{ssr_idx}.to_v_SSR.weight'], + } + ssr_aligner.load_state_dict(attn_dict) + ssr_attentions[('output', i)] = ssr_aligner + ssr_idx += 2 + + ssr_aligner = SSRCrossAttention(ssr_caa_dims['middle'][0]) + attn_dict = { + 'to_k.weight': state_dict[f'{ssr_idx}.to_k_SSR.weight'], + 'to_v.weight': state_dict[f'{ssr_idx}.to_v_SSR.weight'], + } + ssr_aligner.load_state_dict(attn_dict) + ssr_attentions[('middle', 0)] = ssr_aligner + + return (ssr_attentions,) \ No newline at end of file