File size: 2,581 Bytes
82ea528 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
from ..utils.attn_bank import AttentionBank
class LTXAttentionBankNode:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"save_steps": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
"blocks": ("STRING", { "multiline": True })
}}
RETURN_TYPES = ("ATTN_BANK",)
FUNCTION = "build"
CATEGORY = "ltxtricks"
def build(self, save_steps, blocks=''):
block_map = {}
block_list = blocks.split(',')
for block in block_list:
block_idx = int(block)
block_map[block_idx] = {}
bank = AttentionBank(save_steps, block_map)
return (bank, )
class LTXPrepareAttnInjectionsNode:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"latent": ("LATENT",),
"attn_bank": ("ATTN_BANK",),
"query": ("BOOLEAN", { "default": False }),
"key": ("BOOLEAN", { "default": False }),
"value": ("BOOLEAN", { "default": False }),
"inject_steps": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
}, "optional": {
"blocks": ("LTX_BLOCKS",)
}}
RETURN_TYPES = ("LATENT", "ATTN_INJ")
FUNCTION = "prepare"
CATEGORY = "fluxtapoz"
def prepare(self, latent, attn_bank, query, key, value, inject_steps, blocks=None):
if inject_steps > attn_bank['save_steps']:
raise ValueError(f"Can not inject more steps than were saved.")
attn_bank = AttentionBank(attn_bank['save_steps'], attn_bank['block_map'], inject_steps)
attn_bank['inject_settings'] = set([])
if query:
attn_bank['inject_settings'].add('q')
if key:
attn_bank['inject_settings'].add('k')
if value:
attn_bank['inject_settings'].add('v')
if blocks is not None:
attn_bank['block_map'] = {**attn_bank['block_map']}
for key in list(attn_bank['block_map'].keys()):
if key not in blocks:
del attn_bank['block_map'][key]
# Hack to force order of operations in ComfyUI graph
return (latent, attn_bank)
class LTXAttentioOverrideNode:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"blocks": ("STRING", { "multiline": True })
}}
RETURN_TYPES = ("LTX_BLOCKS",)
FUNCTION = "build"
CATEGORY = "ltxtricks"
def build(self, blocks=''):
block_set = set(list(int(block) for block in blocks.split(',')))
return (block_set, )
|