File size: 6,881 Bytes
e545366 | 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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 | """ComfyUI nodes for the Anima style stream (decoupled cross-attention).
Requires a trained AnimaStyleAdapter checkpoint (see style_adapter.py
and the project docs); with an untrained adapter the gates are zero and
the nodes are an exact no-op.
The style stream composes freely with the in-context character stream:
AnimaInContextApply patches self-attention (T-axis reference frames),
AnimaStyleApply patches cross-attention — chain both Apply nodes and
control char/style strength independently.
"""
import os
import torch
import comfy.utils
import folder_paths
from comfy.patcher_extension import WrappersMP
from .style_adapter import AnimaStyleAdapter, StyleState
STYLE_WRAPPER_KEY = "anima_style_ref"
_ADAPTER_DIR = "anima_style_adapters"
if _ADAPTER_DIR not in folder_paths.folder_names_and_paths:
folder_paths.add_model_folder_path(
_ADAPTER_DIR, os.path.join(folder_paths.models_dir, _ADAPTER_DIR)
)
class AnimaStyleAdapterLoader:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"adapter_name": (folder_paths.get_filename_list(_ADAPTER_DIR),),
},
}
RETURN_TYPES = ("ANIMA_STYLE_ADAPTER",)
FUNCTION = "load"
CATEGORY = "anima/style"
def load(self, adapter_name):
path = folder_paths.get_full_path_or_raise(_ADAPTER_DIR, adapter_name)
sd = comfy.utils.load_torch_file(path, safe_load=True)
adapter = AnimaStyleAdapter.from_state_dict(sd)
adapter.eval()
return (adapter,)
class AnimaStyleEncode:
"""Encode style reference image(s) with a SigLIP CLIP_VISION model,
keeping the hidden-state stack the adapter aggregates over.
Multiple images are encoded independently and their patch tokens are
concatenated at apply time (attention pools over all of them)."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"clip_vision": ("CLIP_VISION",),
"image": ("IMAGE",),
},
}
RETURN_TYPES = ("ANIMA_STYLE_EMBEDS",)
FUNCTION = "encode"
CATEGORY = "anima/style"
def encode(self, clip_vision, image):
out = clip_vision.encode_image(image)
hs = out["all_hidden_states"] # (B, L, N, D)
if hs is None:
raise RuntimeError(
"this CLIP_VISION model does not expose all hidden states; "
"use a SigLIP vision model"
)
return ({"hidden_states": hs},)
class AnimaStyleApply:
"""Attach the style stream to an Anima model.
style_weight: global multiplier on the (learned, per-block) gates.
cond_only: apply style only to the cond half of the CFG batch.
start/end_percent: sampling window, same semantics as the
in-context Apply node.
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("MODEL",),
"style_adapter": ("ANIMA_STYLE_ADAPTER",),
"style_embeds": ("ANIMA_STYLE_EMBEDS",),
"style_weight": ("FLOAT", {"default": 1.0, "min": -2.0, "max": 5.0, "step": 0.05}),
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
"end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
},
"optional": {
"cond_only": ("BOOLEAN", {"default": True}),
},
}
RETURN_TYPES = ("MODEL",)
FUNCTION = "apply"
CATEGORY = "anima/style"
def apply(self, model, style_adapter, style_embeds, style_weight,
start_percent, end_percent, cond_only=True):
m = model.clone()
ms = m.get_model_object("model_sampling")
sigma_start = ms.percent_to_sigma(start_percent)
sigma_end = ms.percent_to_sigma(end_percent)
dm = m.get_model_object("diffusion_model")
n_layers = style_adapter.config["n_layers"]
hidden_states = style_embeds["hidden_states"][:, -n_layers:] # (B, K, N, D)
state = StyleState()
# per-(device, dtype) cache of the per-block K/V — the style
# tokens are constant across sampling steps
kv_cache = {}
def compute_kv(device, dtype):
key = (device, dtype)
if key not in kv_cache:
adapter = style_adapter.to(device)
hs = hidden_states.to(device=device, dtype=next(adapter.parameters()).dtype)
with torch.no_grad():
tokens = adapter.aggregator(hs) # (B, N, style_dim)
# multiple style images -> one long token sequence
tokens = tokens.reshape(1, -1, tokens.shape[-1])
kv_cache[key] = [
tuple(t.to(dtype) for t in blk.kv(tokens)) for blk in adapter.blocks
]
return kv_cache[key]
def wrapper(executor, x, timesteps, context, fps=None, padding_mask=None, **kwargs):
to = kwargs.get("transformer_options", {})
sigmas = to.get("sigmas", None)
if sigmas is not None:
s = float(sigmas.max())
if s > sigma_start or s < sigma_end:
return executor(x, timesteps, context, fps, padding_mask, **kwargs)
state.kv_per_block = compute_kv(x.device, x.dtype)
state.weight = style_weight
state.sample_scale_B = None
cou = to.get("cond_or_uncond", None)
if cond_only and cou and x.shape[0] % len(cou) == 0:
chunk = x.shape[0] // len(cou)
scale = torch.ones(x.shape[0], device=x.device)
for i, kind in enumerate(cou):
if kind == 1:
scale[i * chunk:(i + 1) * chunk] = 0.0
state.sample_scale_B = scale
state.active = True
try:
return executor(x, timesteps, context, fps, padding_mask, **kwargs)
finally:
state.active = False
m.add_wrapper_with_key(WrappersMP.DIFFUSION_MODEL, STYLE_WRAPPER_KEY, wrapper)
from .style_adapter import StyleCrossAttention
for i, block in enumerate(dm.blocks):
m.add_object_patch(
"diffusion_model.blocks.{}.cross_attn".format(i),
StyleCrossAttention(block.cross_attn, style_adapter.blocks[i], state, i),
)
return (m,)
NODE_CLASS_MAPPINGS = {
"AnimaStyleAdapterLoader": AnimaStyleAdapterLoader,
"AnimaStyleEncode": AnimaStyleEncode,
"AnimaStyleApply": AnimaStyleApply,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"AnimaStyleAdapterLoader": "Anima Style Adapter Loader",
"AnimaStyleEncode": "Anima Style Encode (SigLIP)",
"AnimaStyleApply": "Anima Style Apply",
}
|