93 lines
4.1 KiB
Python
93 lines
4.1 KiB
Python
import torch
|
|
from comfy.ldm.modules.attention import optimized_attention
|
|
from ... import ROOT_NAME, SYMBOL, NODE_SURFIX
|
|
from comfy_api.v0_0_2 import io
|
|
|
|
def attention_pytorch(q, k, v, heads, temperature=1.0, mask=None):
|
|
b, _, dim_head = q.shape
|
|
dim_head //= heads
|
|
q, k, v = map(
|
|
lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2),
|
|
(q, k, v),
|
|
)
|
|
|
|
scale = (dim_head ** -0.5) / temperature
|
|
|
|
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 AttentionScale(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id=f"AttentionScale{NODE_SURFIX}",
|
|
display_name=f"Attention Scale {SYMBOL}",
|
|
category=ROOT_NAME + "for_test",
|
|
inputs=[
|
|
io.Model.Input("model"),
|
|
io.Float.Input("temperature", default=1.0, min=-1000.0, max=1000.0, step=0.01),
|
|
io.Float.Input("start_step", default=0, min=0, max=1, step=0.001),
|
|
io.Float.Input("end_step", default=1, min=0, max=1, step=0.001),
|
|
io.Boolean.Input("attn1", default=True),
|
|
io.Boolean.Input("attn2", default=True),
|
|
],
|
|
outputs=[
|
|
io.Model.Output(),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model, temperature, start_step, end_step, attn1, attn2) -> io.NodeOutput:
|
|
new_model = model.clone()
|
|
|
|
temperature_ = temperature
|
|
start_sigma = new_model.model.model_sampling.percent_to_sigma(start_step)
|
|
end_sigma = new_model.model.model_sampling.percent_to_sigma(end_step)
|
|
|
|
def attn_patch(q, k, v, extra_options):
|
|
sigma = extra_options["sigmas"][0].item()
|
|
|
|
if end_sigma <= sigma <= start_sigma:
|
|
output = attention_pytorch(q, k, v, extra_options["n_heads"], temperature = temperature_)
|
|
else:
|
|
output = attention_pytorch(q, k, v, extra_options["n_heads"], temperature = 1.0)
|
|
|
|
return output
|
|
|
|
def dummy_attn_path(q, k, v, extra_options):
|
|
return optimized_attention(q, k, v, extra_options["n_heads"])
|
|
|
|
sdxl = hasattr(new_model.model.diffusion_model, "label_emb")
|
|
|
|
attn1_patch = attn_patch if attn1 else dummy_attn_path
|
|
attn2_patch = attn_patch if attn2 else dummy_attn_path
|
|
|
|
if not sdxl:
|
|
for id in [1,2,4,5,7,8]: # id of input_blocks that have cross attention
|
|
new_model.set_model_attn1_replace(attn1_patch, "input", id)
|
|
new_model.set_model_attn2_replace(attn2_patch, "input", id)
|
|
new_model.set_model_attn1_replace(attn1_patch, "middle", 0)
|
|
new_model.set_model_attn2_replace(attn2_patch, "middle", 0)
|
|
for id in [3,4,5,6,7,8,9,10,11]: # id of output_blocks that have cross attention
|
|
new_model.set_model_attn1_replace(attn1_patch, "output", id)
|
|
new_model.set_model_attn2_replace(attn2_patch, "output", id)
|
|
else:
|
|
for id in [4,5,7,8]: # id of input_blocks that have cross attention
|
|
block_indices = range(2) if id in [4, 5] else range(10) # transformer_depth
|
|
for index in block_indices:
|
|
new_model.set_model_attn1_replace(attn1_patch, "input", id, index)
|
|
new_model.set_model_attn2_replace(attn2_patch, "input", id, index)
|
|
for index in range(10):
|
|
new_model.set_model_attn1_replace(attn1_patch, "middle", 0, index)
|
|
new_model.set_model_attn2_replace(attn2_patch, "middle", 0, index)
|
|
for id in range(6): # id of output_blocks that have cross attention
|
|
block_indices = range(2) if id in [3, 4, 5] else range(10) # transformer_depth
|
|
for index in block_indices:
|
|
new_model.set_model_attn1_replace(attn1_patch, "output", id, index)
|
|
new_model.set_model_attn2_replace(attn2_patch, "output", id, index)
|
|
|
|
return io.NodeOutput(new_model)
|