370 lines
15 KiB
Python
370 lines
15 KiB
Python
"""ComfyUI nodes for the native MS-Diffusion chain.
|
|
|
|
Intended graph (everything except the three MS nodes is a stock ComfyUI node):
|
|
|
|
CheckpointLoaderSimple -> MODEL / CLIP / VAE
|
|
LoadImage xN -> native ImageBatch (chained N-1 times) -> CLIPVisionEncode
|
|
-> CLIP_VISION_OUTPUT
|
|
CLIPTextEncode x2 -> positive / negative CONDITIONING
|
|
MSDiffusionLoader ------------------------------> MS_ADAPTER
|
|
MSDiffusionConditioning ------------------------> MS_EMBEDS (+ clean prompt)
|
|
MSDiffusionApply(MODEL, MS_ADAPTER, MS_EMBEDS) --> patched MODEL
|
|
EmptyLatentImage -> KSampler -> VAEDecode -> SaveImage
|
|
|
|
No diffusers / transformers / omegaconf / einops / PIL imports anywhere.
|
|
"""
|
|
|
|
import os
|
|
import re
|
|
|
|
import torch
|
|
import folder_paths
|
|
|
|
from .core import (MSAdapter, MSEmbeds, adapter_hidden_sizes, build_block_weights,
|
|
enumerate_attn2, make_attn2_patch, new_state, NUM_TOKENS,
|
|
NUM_DUMMY_TOKENS)
|
|
|
|
ADAPTER_FOLDER = "ms_adapter"
|
|
_ADAPTER_EXTS = (".bin", ".safetensors", ".pt", ".pth", ".ckpt")
|
|
|
|
|
|
def _register_folders():
|
|
"""Register models/ipadapter and models/photomaker as one virtual folder.
|
|
|
|
The upstream plugin hard-coded models/photomaker; the user's actual file is
|
|
in models/ipadapter. Support both without touching ComfyUI source.
|
|
"""
|
|
try:
|
|
if ADAPTER_FOLDER not in folder_paths.folder_names_and_paths:
|
|
folder_paths.folder_names_and_paths[ADAPTER_FOLDER] = ([], set())
|
|
for sub in ("ipadapter", "photomaker"):
|
|
p = os.path.join(folder_paths.models_dir, sub)
|
|
if os.path.isdir(p):
|
|
folder_paths.add_model_folder_path(ADAPTER_FOLDER, p)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
_register_folders()
|
|
|
|
|
|
def _adapter_files():
|
|
try:
|
|
files = [f for f in folder_paths.get_filename_list(ADAPTER_FOLDER)
|
|
if f.lower().endswith(_ADAPTER_EXTS)]
|
|
except Exception:
|
|
files = []
|
|
if not files:
|
|
# fallback scan, keeps the node usable even if registration was blocked
|
|
for sub in ("ipadapter", "photomaker"):
|
|
d = os.path.join(folder_paths.models_dir, sub)
|
|
if os.path.isdir(d):
|
|
files += [sub + "/" + f for f in sorted(os.listdir(d))
|
|
if f.lower().endswith(_ADAPTER_EXTS)]
|
|
files = sorted(set(files)) or ["none"]
|
|
# models/ipadapter often holds many unrelated IP-Adapter weights; put the
|
|
# real MS-Diffusion checkpoint first so the default is the correct file.
|
|
preferred = [f for f in files if os.path.basename(f).lower().startswith("ms_adapter")]
|
|
return preferred + [f for f in files if f not in preferred]
|
|
|
|
|
|
def _adapter_path(name):
|
|
try:
|
|
return folder_paths.get_full_path_or_raise(ADAPTER_FOLDER, name)
|
|
except Exception:
|
|
return os.path.join(folder_paths.models_dir, name.replace("/", os.sep))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# phrase token position (port of msdiffusion/utils.py, no transformers import)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _token_encoder(clip):
|
|
"""Return f(word) -> len(HF tokenizer.encode(word)).
|
|
|
|
ComfyUI already owns a CLIP tokenizer instance; we just borrow it, so this
|
|
plugin adds no transformers import of its own.
|
|
"""
|
|
try:
|
|
tk = clip.tokenizer.clip_l.tokenizer
|
|
return lambda w: len(tk.encode(w))
|
|
except Exception:
|
|
pass
|
|
# fallback: derive from ComfyUI's own token/weight pairs
|
|
def from_comfy(w):
|
|
pairs = clip.tokenize(w)["l"][0]
|
|
n = 0
|
|
for tid, _ in pairs:
|
|
if int(tid) == 49407: # EOS == pad token
|
|
break
|
|
n += 1
|
|
return max(n, 1) + 1 # [BOS] + word tokens => n-1 + 2
|
|
return from_comfy
|
|
|
|
|
|
def _is_equal_words(pr_words, ph_words):
|
|
if len(pr_words) != len(ph_words):
|
|
return False
|
|
for pr, ph in zip(pr_words, ph_words):
|
|
if "-" + ph not in pr and ph != re.sub(r'[.!?,:]$', '', pr):
|
|
return False
|
|
return True
|
|
|
|
|
|
def _get_phrase_idx(enc_len, phrase, prompt, num=0):
|
|
phrase_words = phrase.split()
|
|
if len(phrase_words) == 0:
|
|
return [0, 0]
|
|
prompt_words = prompt.split()
|
|
start, end = 1, 0
|
|
for i in range(len(prompt_words)):
|
|
if _is_equal_words(prompt_words[i:i + len(phrase_words)], phrase_words):
|
|
if num != 0:
|
|
num -= 1
|
|
continue
|
|
end = start
|
|
prompt_words[i + len(phrase_words) - 1] = re.sub(
|
|
r'[.!?,:]$', '', prompt_words[i + len(phrase_words) - 1])
|
|
for j in range(i, i + len(phrase_words)):
|
|
end += enc_len(prompt_words[j]) - 2
|
|
break
|
|
else:
|
|
start += enc_len(prompt_words[i]) - 2
|
|
if end == 0:
|
|
return [0, 0]
|
|
return [start, end]
|
|
|
|
|
|
def _parse_boxes(text, num_ref, enabled):
|
|
"""Upstream: mask_threshold == 0 disables layout guidance -> all-zero boxes."""
|
|
zeros = [[0.0, 0.0, 0.0, 0.0] for _ in range(num_ref)]
|
|
if not enabled or not text.strip():
|
|
return zeros
|
|
groups = [g.strip() for g in text.strip().split(";") if g.strip()]
|
|
parsed = []
|
|
for g in groups:
|
|
vals = [float(x.strip()) for x in g.split(",")]
|
|
if len(vals) != 4:
|
|
raise ValueError("MS-Diffusion: each box must have 4 floats (x1,y1,x2,y2); got %r" % g)
|
|
parsed.append(vals)
|
|
if len(parsed) < num_ref:
|
|
# A zero box masks that subject's 16 image tokens completely (-10000 bias),
|
|
# silently dropping it. Upstream raises here instead of guessing; do the same.
|
|
raise ValueError(
|
|
"MS-Diffusion: %d subjects but only %d boxes. Provide one "
|
|
"'x1,y1,x2,y2' per subject separated by ';' (or set mask_threshold=0 "
|
|
"to disable layout guidance)." % (num_ref, len(parsed)))
|
|
return parsed[:num_ref]
|
|
|
|
|
|
def _clip_l_pooled(clip, phrase):
|
|
"""CLIP-L 768-d pooled output (the grounding phrase embedding).
|
|
|
|
SDXL's CONDITIONING pooled_output is clip_g (1280) and cannot be used here.
|
|
"""
|
|
clip.load_model()
|
|
dev = clip.patcher.load_device
|
|
try:
|
|
clip.cond_stage_model.clip_l.set_clip_options({"execution_device": dev})
|
|
except Exception:
|
|
pass
|
|
tokens = clip.tokenize(phrase)["l"]
|
|
with torch.no_grad():
|
|
out = clip.cond_stage_model.clip_l.encode_token_weights(tokens)
|
|
pooled = out[1]
|
|
return pooled.to(torch.float32).cpu()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Nodes
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class MSDiffusionLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"ms_adapter": (_adapter_files(),)}}
|
|
|
|
RETURN_TYPES = ("MS_ADAPTER",)
|
|
RETURN_NAMES = ("ms_adapter",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "MSdiffusion/Native"
|
|
DESCRIPTION = ("Load ms_adapter.bin from models/ipadapter or models/photomaker. "
|
|
"Pure torch; no diffusers/transformers needed.")
|
|
|
|
def load(self, ms_adapter):
|
|
if ms_adapter == "none":
|
|
raise RuntimeError("MS-Diffusion: no ms_adapter.bin found in models/ipadapter or models/photomaker")
|
|
path = _adapter_path(ms_adapter)
|
|
sizes = adapter_hidden_sizes(path)
|
|
adapter = MSAdapter(sizes).load(path)
|
|
print("[MS-Diffusion] loaded %s (%d cross-attn groups)" % (path, len(sizes)))
|
|
return (adapter,)
|
|
|
|
|
|
class MSDiffusionConditioning:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {
|
|
"ms_adapter": ("MS_ADAPTER",),
|
|
"clip": ("CLIP",),
|
|
"clip_vision_output": ("CLIP_VISION_OUTPUT",),
|
|
"prompt": ("STRING", {"multiline": True,
|
|
"default": "best quality, high quality, a [dog] on the beach"}),
|
|
"boxes": ("STRING", {"multiline": True,
|
|
"default": "0.25, 0.25, 0.75, 0.75",
|
|
"tooltip": "one 'x1,y1,x2,y2' per subject, separated by ';'. "
|
|
"Ignored when mask_threshold is 0."}),
|
|
"mask_threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05,
|
|
"tooltip": "0 disables layout guidance entirely (upstream behaviour). "
|
|
">0 enables box guidance + pseudo attention mask."}),
|
|
"drop_grounding_tokens": ("BOOLEAN", {"default": False}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("MS_EMBEDS", "STRING")
|
|
RETURN_NAMES = ("ms_embeds", "prompt_clean")
|
|
FUNCTION = "encode"
|
|
CATEGORY = "MSdiffusion/Native"
|
|
DESCRIPTION = ("Build MS-Diffusion image prompt tokens from CLIP-Vision penultimate "
|
|
"hidden states + CLIP-L phrase pooled embeddings + layout boxes. "
|
|
"Connect prompt_clean to your positive CLIPTextEncode so the two match.")
|
|
|
|
def encode(self, ms_adapter, clip, clip_vision_output, prompt, boxes,
|
|
mask_threshold, drop_grounding_tokens):
|
|
import comfy.model_management as mm
|
|
|
|
phrases = [p.strip() for p in re.findall(r'\[(.*?)\]', prompt)]
|
|
if not phrases:
|
|
raise ValueError("MS-Diffusion: subject names must be enclosed in [ ] inside the prompt")
|
|
for p in phrases:
|
|
if " " in p:
|
|
raise ValueError("MS-Diffusion: subject %r must be a single word" % p)
|
|
phrases = list(dict.fromkeys(phrases))
|
|
clean_prompt = prompt.replace("]", " ").replace("[", " ")
|
|
|
|
hidden = clip_vision_output["penultimate_hidden_states"]
|
|
if hidden.dim() != 3:
|
|
raise ValueError("MS-Diffusion: CLIP-Vision output must be penultimate_hidden_states [B,257,1664]")
|
|
num_ref = hidden.shape[0]
|
|
if num_ref != len(phrases):
|
|
raise ValueError(
|
|
"MS-Diffusion: got %d subject images but %d [subject] names in prompt (%s)"
|
|
% (num_ref, len(phrases), phrases))
|
|
if hidden.shape[-1] != 1664:
|
|
raise ValueError("MS-Diffusion: expected CLIP-ViT-bigG hidden size 1664, got %d. "
|
|
"Use clip_vision_g.safetensors." % hidden.shape[-1])
|
|
|
|
layout_on = mask_threshold > 0.0
|
|
box_vals = _parse_boxes(boxes, num_ref, layout_on)
|
|
boxes_t = torch.tensor([box_vals], dtype=torch.float32) # [1, num_ref, 4]
|
|
batch_boxes = boxes_t.view(num_ref, 4) # Resampler grounding input
|
|
|
|
phrase_embeds = torch.cat([_clip_l_pooled(clip, p) for p in phrases], dim=0)
|
|
if phrase_embeds.shape != (num_ref, 768):
|
|
raise RuntimeError("MS-Diffusion: CLIP-L pooled shape %s, expected (%d,768)"
|
|
% (tuple(phrase_embeds.shape), num_ref))
|
|
|
|
enc_len = _token_encoder(clip)
|
|
seen = {}
|
|
pidx = []
|
|
for p in phrases:
|
|
n = seen.get(p, 0)
|
|
seen[p] = n + 1
|
|
pidx.append(_get_phrase_idx(enc_len, p, clean_prompt, num=n))
|
|
phrase_idxes = torch.tensor([pidx], dtype=torch.int32) # [1, num_ref, 2]
|
|
|
|
dev = mm.get_torch_device()
|
|
dtype = torch.float16
|
|
with torch.no_grad():
|
|
ms_adapter.image_proj.to(device=dev, dtype=dtype)
|
|
x = hidden.to(device=dev, dtype=dtype)
|
|
grounding = {
|
|
"boxes": batch_boxes.to(device=dev, dtype=dtype),
|
|
"phrase_embeds": phrase_embeds.to(device=dev, dtype=dtype),
|
|
"drop_grounding_tokens": [1 if drop_grounding_tokens else 0],
|
|
}
|
|
ip_tokens = ms_adapter.image_proj(x, grounding_kwargs=grounding)
|
|
ip_tokens = ip_tokens.view(1, -1, ip_tokens.shape[-2], ip_tokens.shape[-1])
|
|
ip_tokens = ip_tokens.view(1, ip_tokens.shape[-3] * ip_tokens.shape[-2], ip_tokens.shape[-1])
|
|
dummy = ms_adapter.dummy_image_tokens.to(device=dev, dtype=dtype)
|
|
image_prompt_embeds = torch.cat([dummy, ip_tokens], dim=1).float().cpu()
|
|
ms_adapter.image_proj.to(device=torch.device("cpu"))
|
|
ms_adapter._placed = None
|
|
|
|
assert image_prompt_embeds.shape == (1, NUM_DUMMY_TOKENS + NUM_TOKENS * num_ref, 2048)
|
|
embeds = MSEmbeds(
|
|
image_prompt_embeds=image_prompt_embeds,
|
|
boxes=boxes_t,
|
|
phrase_idxes=phrase_idxes,
|
|
num_ref=num_ref,
|
|
scale=1.0, # set by MSDiffusionApply
|
|
mask_threshold=None if not layout_on else float(mask_threshold),
|
|
start_step=5, # set by MSDiffusionApply
|
|
)
|
|
return (embeds, clean_prompt)
|
|
|
|
|
|
class MSDiffusionApply:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {
|
|
"model": ("MODEL",),
|
|
"ms_adapter": ("MS_ADAPTER",),
|
|
"ms_embeds": ("MS_EMBEDS",),
|
|
"scale": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 2.0, "step": 0.05}),
|
|
"start_step": ("INT", {"default": 5, "min": 0, "max": 1024,
|
|
"tooltip": "pseudo attention mask kicks in after this many "
|
|
"denoising steps (only when mask_threshold > 0)"}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "apply"
|
|
CATEGORY = "MSdiffusion/Native"
|
|
DESCRIPTION = ("Patch every native SDXL attn2 with MS-Diffusion dual-branch "
|
|
"cross-attention. Output is a normal MODEL: use stock KSampler / "
|
|
"ControlNet / FreeU / LoRA nodes downstream.")
|
|
|
|
def apply(self, model, ms_adapter, ms_embeds, scale, start_step):
|
|
import copy
|
|
m = model.clone()
|
|
desc = enumerate_attn2(m.model.diffusion_model)
|
|
if not desc:
|
|
raise RuntimeError("MS-Diffusion: no attn2 modules found. This adapter requires a stock SDXL UNet.")
|
|
perm = build_block_weights(desc, ms_adapter.hidden_sizes)
|
|
|
|
emb = copy.copy(ms_embeds)
|
|
emb.scale = float(scale)
|
|
emb.start_step = int(start_step)
|
|
sid = new_state(ms_adapter, emb)
|
|
|
|
# IMPORTANT: set_model_patch_replace() REPLACES the whole
|
|
# transformer_options dict and silently drops previously attached keys.
|
|
# Build patches_replace in place on the cloned options instead.
|
|
to = m.model_options["transformer_options"]
|
|
pr = dict(to.get("patches_replace", {}))
|
|
a2 = dict(pr.get("attn2", {}))
|
|
for pos, (key3, _inner, _section) in enumerate(desc):
|
|
a2[key3] = make_attn2_patch(sid, perm[pos], is_sentinel=(pos == 0))
|
|
pr["attn2"] = a2
|
|
to["patches_replace"] = pr
|
|
print("[MS-Diffusion] patched %d attn2 layers (scale=%.2f, layout=%s)"
|
|
% (len(desc), emb.scale, "on" if emb.mask_threshold is not None else "off"))
|
|
return (m,)
|
|
|
|
|
|
# No custom image-batch node: ComfyUI's native `ImageBatch` already stacks
|
|
# images into [N,H,W,3] and aligns size/channels. Chain it for N > 2.
|
|
# (The old `MS_Object_img_Batch` / `MSDiffusionImageBatch` ids are intentionally
|
|
# gone; replace them with the native ImageBatch node in old graphs.)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"MSDiffusionLoader": MSDiffusionLoader,
|
|
"MSDiffusionConditioning": MSDiffusionConditioning,
|
|
"MSDiffusionApply": MSDiffusionApply,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"MSDiffusionLoader": "MS-Diffusion Loader (Native)",
|
|
"MSDiffusionConditioning": "MS-Diffusion Conditioning (Native)",
|
|
"MSDiffusionApply": "MS-Diffusion Apply to MODEL (Native)",
|
|
}
|