224 lines
8.5 KiB
Python
224 lines
8.5 KiB
Python
import torch
|
|
from comfy_api.v0_0_2 import io
|
|
from ... import ROOT_NAME, SYMBOL, NODE_SURFIX
|
|
|
|
CATEGORY_NAME = ROOT_NAME + "reference"
|
|
|
|
class ReferenceApply(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id=f"ReferenceApply{NODE_SURFIX}",
|
|
display_name=f"Reference Apply {SYMBOL}",
|
|
category=CATEGORY_NAME,
|
|
inputs=[
|
|
io.Model.Input("model"),
|
|
io.Int.Input("index", default=0, min=0, max=256),
|
|
io.Combo.Input("mode", options=["concat", "replace"], default="concat"),
|
|
io.Int.Input("depth", default=12, min=-1, max=12),
|
|
io.Float.Input("start_step", default=0, min=0, max=1, step=0.01),
|
|
io.Float.Input("end_step", default=1, min=0, max=1, step=0.01),
|
|
io.Boolean.Input("apply_input", default=True),
|
|
io.Boolean.Input("apply_middle", default=True),
|
|
io.Boolean.Input("apply_output", default=True),
|
|
],
|
|
outputs=[
|
|
io.Model.Output(),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model, index, mode, depth, start_step, end_step, apply_input, apply_middle, apply_output) -> io.NodeOutput:
|
|
model_reference = model.clone()
|
|
start_sigma = model_reference.model.model_sampling.percent_to_sigma(start_step)
|
|
end_sigma = model_reference.model.model_sampling.percent_to_sigma(end_step)
|
|
|
|
sdxl = hasattr(model_reference.model.diffusion_model, "label_emb")
|
|
num_blocks = 8 if sdxl else 11
|
|
|
|
def reference_apply(q, k, v, extra_options):
|
|
block_name, block_id = extra_options["block"]
|
|
chunks = len(extra_options["cond_or_uncond"])
|
|
batch_size = q.shape[0] // chunks
|
|
|
|
if block_name == "input" and not apply_input:
|
|
return q, k, v
|
|
if block_name == "middle" and not apply_middle:
|
|
return q, k, v
|
|
if block_name == "output" and not apply_output:
|
|
return q, k, v
|
|
|
|
if block_name == "output":
|
|
block_number = num_blocks - block_id
|
|
else:
|
|
block_number = block_id
|
|
|
|
q_out = q.clone()
|
|
k_out = k.clone()
|
|
v_out = v.clone()
|
|
|
|
sigma = extra_options["sigmas"][0].item()
|
|
|
|
if end_sigma <= sigma <= start_sigma and block_number <= depth:
|
|
k_ref = k_out[index::batch_size].repeat_interleave(batch_size, dim=0).clone()
|
|
v_ref = v_out[index::batch_size].repeat_interleave(batch_size, dim=0).clone()
|
|
|
|
k_out = torch.cat([k_out, k_ref], dim=1) if mode == "concat" else k_ref
|
|
v_out = torch.cat([v_out, v_ref], dim=1) if mode == "concat" else v_ref
|
|
|
|
return q_out, k_out, v_out
|
|
|
|
model_reference.set_model_attn1_patch(reference_apply)
|
|
|
|
return io.NodeOutput(model_reference)
|
|
|
|
class ReferenceLatent(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id=f"ReferenceLatent{NODE_SURFIX}",
|
|
display_name=f"Reference Latent {SYMBOL}",
|
|
category=CATEGORY_NAME,
|
|
inputs=[
|
|
io.Latent.Input("latent"),
|
|
io.Int.Input("index", default=0, min=0, max=256),
|
|
io.Int.Input("batch_size", default=1, min=1, max=256),
|
|
],
|
|
outputs=[
|
|
io.Latent.Output(),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, latent, index, batch_size) -> io.NodeOutput:
|
|
latent_new = latent.copy()
|
|
|
|
sample = latent_new["samples"]
|
|
height, width = sample.shape[2], sample.shape[3]
|
|
|
|
empty_latent = torch.zeros_like(latent["samples"]).repeat(batch_size , 1, 1, 1)
|
|
empty_latent[index] = sample[0]
|
|
noise_mask = torch.ones(batch_size, 1, height * 8, width * 8).to(sample)
|
|
noise_mask[index] = 0.0
|
|
|
|
latent_new["samples"] = empty_latent
|
|
latent_new["noise_mask"] = noise_mask
|
|
|
|
return io.NodeOutput(latent_new)
|
|
|
|
class MultipleReferenceApply(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id=f"MultipleReferenceApply{NODE_SURFIX}",
|
|
display_name=f"Multiple Reference Apply {SYMBOL}",
|
|
category=CATEGORY_NAME,
|
|
inputs=[
|
|
io.Model.Input("model"),
|
|
io.String.Input("indices", default="0"),
|
|
io.Int.Input("depth", default=12, min=-1, max=12),
|
|
io.Float.Input("start_step", default=0, min=0, max=1, step=0.01),
|
|
io.Float.Input("end_step", default=1, min=0, max=1, step=0.01),
|
|
io.Boolean.Input("apply_input", default=True),
|
|
io.Boolean.Input("apply_middle", default=True),
|
|
io.Boolean.Input("apply_output", default=True),
|
|
io.String.Input("weights", default=""),
|
|
],
|
|
outputs=[
|
|
io.Model.Output(),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model, indices, depth, start_step, end_step, apply_input, apply_middle, apply_output, weights) -> io.NodeOutput:
|
|
model_reference = model.clone()
|
|
start_sigma = model_reference.model.model_sampling.percent_to_sigma(start_step)
|
|
end_sigma = model_reference.model.model_sampling.percent_to_sigma(end_step)
|
|
|
|
sdxl = hasattr(model_reference.model.diffusion_model, "label_emb")
|
|
num_blocks = 8 if sdxl else 11
|
|
|
|
indices = [int(i) for i in indices.split(",") if i.strip().isdigit()]
|
|
weights = [float(i) for i in weights.split(",") if i.strip()] if weights else [1.0] * len(indices)
|
|
|
|
def reference_apply(q, k, v, extra_options):
|
|
block_name, block_id = extra_options["block"]
|
|
|
|
|
|
if block_name == "input" and not apply_input:
|
|
return q, k, v
|
|
if block_name == "middle" and not apply_middle:
|
|
return q, k, v
|
|
if block_name == "output" and not apply_output:
|
|
return q, k, v
|
|
|
|
if block_name == "output":
|
|
block_number = num_blocks - block_id
|
|
else:
|
|
block_number = block_id
|
|
|
|
q_out = q.clone()
|
|
k_out = k.clone()
|
|
v_out = v.clone()
|
|
|
|
sigma = extra_options["sigmas"][0].item()
|
|
|
|
|
|
if end_sigma <= sigma <= start_sigma and block_number <= depth:
|
|
chunks = len(extra_options["cond_or_uncond"])
|
|
batch_size = q.shape[0] // chunks
|
|
num_tokens = q.shape[1]
|
|
|
|
k_refs = torch.cat([k_out[i::batch_size] for i in indices], dim=1)
|
|
v_refs = torch.cat([v_out[i::batch_size] * weight for i, weight in zip(indices, weights)], dim=1)
|
|
|
|
k_out = k_out.repeat(1, len(indices)+1, 1).clone()
|
|
v_out = v_out.repeat(1, len(indices)+1, 1).clone()
|
|
for i in range(batch_size):
|
|
if i not in indices:
|
|
k_out[i::batch_size, num_tokens:] = k_refs.clone()
|
|
v_out[i::batch_size, num_tokens:] = v_refs.clone()
|
|
|
|
return q_out, k_out, v_out
|
|
|
|
model_reference.set_model_attn1_patch(reference_apply)
|
|
|
|
return io.NodeOutput(model_reference)
|
|
|
|
class MultipleReferenceLatent(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id=f"MultipleReferenceLatent{NODE_SURFIX}",
|
|
display_name=f"Multiple Reference Latent {SYMBOL}",
|
|
category=CATEGORY_NAME,
|
|
inputs=[
|
|
io.Latent.Input("latent"),
|
|
io.String.Input("indices", default="0"),
|
|
io.Int.Input("batch_size", default=1, min=1, max=256),
|
|
],
|
|
outputs=[
|
|
io.Latent.Output(),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, latent, indices, batch_size) -> io.NodeOutput:
|
|
latent_new = latent.copy()
|
|
indices = [int(i) for i in indices.split(",") if i.strip().isdigit()]
|
|
|
|
sample = latent_new["samples"]
|
|
b, _, height, width = sample.shape
|
|
|
|
assert len(indices) == b
|
|
|
|
empty_latent = torch.zeros_like(latent["samples"][:1]).repeat(batch_size , 1, 1, 1)
|
|
empty_latent[torch.tensor(indices)] = sample
|
|
noise_mask = torch.ones(batch_size, 1, height * 8, width * 8).to(sample)
|
|
noise_mask[torch.tensor(indices)] = 0.0
|
|
|
|
latent_new["samples"] = empty_latent
|
|
latent_new["noise_mask"] = noise_mask
|
|
|
|
return io.NodeOutput(latent_new)
|