nodes/smthemex-ComfyUI_MS_Diffusion
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

init

Browse Source
This commit is contained in:
smthemex
2026-11-06 13:47:26 +08:00
parent a5b97bb7ce
commit 343b27520a
45 changed files with 1822 additions and 199653 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -0,0 +1 @@
__pycache__/
MS_Diffusion.py
-648
View File
@@ -1,648 +0,0 @@
# !/usr/bin/env python
# -*- coding: UTF-8 -*-
import numpy as np
import torch
import os
import sys
import re
import yaml
from PIL import Image
from huggingface_hub import hf_hub_download
from diffusers import (StableDiffusionXLPipeline, DDIMScheduler, ControlNetModel,
KDPM2AncestralDiscreteScheduler, LMSDiscreteScheduler,
DPMSolverMultistepScheduler, DPMSolverSinglestepScheduler,
EulerDiscreteScheduler, HeunDiscreteScheduler,
KDPM2DiscreteScheduler,
EulerAncestralDiscreteScheduler, UniPCMultistepScheduler, AutoencoderKL,
StableDiffusionXLControlNetPipeline, DDPMScheduler, LCMScheduler)
from transformers import CLIPVisionModelWithProjection
from transformers import CLIPImageProcessor
from .msdiffusion.models.projection import Resampler
from .msdiffusion.models.model import MSAdapter
from .msdiffusion.utils import get_phrase_idx, get_eot_idx
from diffusers.schedulers.scheduling_ddim import DDIMScheduler
import folder_paths
from comfy.utils import common_upscale
from omegaconf import OmegaConf
MAX_SEED = np.iinfo(np.int32).max
dir_path = os.path.dirname(os.path.abspath(__file__))
path_dir = os.path.dirname(dir_path)
file_path = os.path.dirname(path_dir)
def get_instance_path(path):
instance_path = os.path.normpath(path)
if sys.platform == 'win32':
instance_path = instance_path.replace('\\', "/")
return instance_path
loras_path = get_instance_path(os.path.join(dir_path, "config", "lora.yaml"))
def get_lora_dict():
# 打开并读取YAML文件
with open(loras_path, 'r', encoding="UTF-8") as stream:
try:
# 解析YAML文件内容
data = yaml.safe_load(stream)
# 此时 'data' 是一个Python字典，里面包含了YAML文件的所有数据
# print(data)
return data
except yaml.YAMLError as exc:
# 如果在解析过程中发生了错误，打印异常信息
print(exc)
datas = get_lora_dict()
lora_lightning_list = datas["lightning_xl_lora"]
scheduler_list = [
"Euler", "Euler a", "DDIM", "DDPM", "DPM++ 2M", "DPM++ 2M Karras", "DPM++ 2M SDE", "DPM++ 2M SDE Karras",
"DPM++ SDE", "DPM++ SDE Karras", "DPM2",
"DPM2 Karras", "DPM2 a", "DPM2 a Karras", "Heun", "LCM", "LMS", "LMS Karras", "UniPC"
]
def phi2narry(img):
img = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)
return img
def nomarl_upscale(img_tensor, width, height):
samples = img_tensor.movedim(-1, 1)
img = common_upscale(samples, width, height, "nearest-exact", "center")
samples = img.movedim(1, -1)
return samples
def phi2tensor(img):
image = np.array(img).astype(np.float32) / 255.0
image = torch.from_numpy(image)[None,]
return image
def add_pil(list, list_add, num):
new_list = list[:num] + list_add + list[num:]
return new_list
def tensor_to_image(tensor):
# tensor = tensor.cpu()
image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
image = Image.fromarray(image_np, mode='RGB')
return image
# get fonts list
def has_parentheses(s):
return bool(re.search(r'\(.*?\)', s))
def contains_brackets(s):
return '[' in s or ']' in s
def narry_list(list_in):
for i in range(len(list_in)):
value = list_in[i]
modified_value = phi2narry(value)
list_in[i] = modified_value
return list_in
def phi_list(list_in):
for i in range(len(list_in)):
value = list_in[i]
list_in[i] = value
return list_in
def narry_list_pil(list_in):
for i in range(len(list_in)):
value = list_in[i]
modified_value = tensor_to_image(value)
list_in[i] = modified_value
return list_in
def get_local_path(file_path, model_path, file):
path = os.path.join(file_path, "models", file, model_path)
model_path = os.path.normpath(path)
if sys.platform.startswith('win32'):
model_path = model_path.replace('\\', "/")
return model_path
def extract_content_between_brackets(text):
# 正则表达式匹配多对括号内的内容
return re.findall(r'\((.*?)\)', text)
def extract_content_from_brackets(text):
# 正则表达式匹配多对方括号内的内容
return re.findall(r'\[(.*?)\]', text)
def get_image_path_list(folder_name):
image_basename_list = os.listdir(folder_name)
image_path_list = sorted(
[os.path.join(folder_name, basename) for basename in image_basename_list]
)
return image_path_list
def get_scheduler(name):
scheduler = False
if name == "Euler":
scheduler = EulerDiscreteScheduler()
elif name == "Euler a":
scheduler = EulerAncestralDiscreteScheduler()
elif name == "DDIM":
scheduler = DDIMScheduler()
elif name == "DDPM":
scheduler = DDPMScheduler()
elif name == "DPM++ 2M":
scheduler = DPMSolverMultistepScheduler()
elif name == "DPM++ 2M Karras":
scheduler = DPMSolverMultistepScheduler(use_karras_sigmas=True)
elif name == "DPM++ 2M SDE":
scheduler = DPMSolverMultistepScheduler(algorithm_type="sde-dpmsolver++")
elif name == "DPM++ 2M SDE Karras":
scheduler = DPMSolverMultistepScheduler(use_karras_sigmas=True, algorithm_type="sde-dpmsolver++")
elif name == "DPM++ SDE":
scheduler = DPMSolverSinglestepScheduler()
elif name == "DPM++ SDE Karras":
scheduler = DPMSolverSinglestepScheduler(use_karras_sigmas=True)
elif name == "DPM2":
scheduler = KDPM2DiscreteScheduler()
elif name == "DPM2 Karras":
scheduler = KDPM2DiscreteScheduler(use_karras_sigmas=True)
elif name == "DPM2 a":
scheduler = KDPM2AncestralDiscreteScheduler()
elif name == "DPM2 a Karras":
scheduler = KDPM2AncestralDiscreteScheduler(use_karras_sigmas=True)
elif name == "Heun":
scheduler = HeunDiscreteScheduler()
elif name == "LCM":
scheduler = LCMScheduler()
elif name == "LMS":
scheduler = LMSDiscreteScheduler()
elif name == "LMS Karras":
scheduler = LMSDiscreteScheduler(use_karras_sigmas=True)
elif name == "UniPC":
scheduler = UniPCMultistepScheduler()
return scheduler
def nomarl_upscale_topil(img_tensor, width, height):
samples = img_tensor.movedim(-1, 1)
img = common_upscale(samples, width, height, "nearest-exact", "center")
samples = img.movedim(1, -1)
img_pil = tensor_to_image(samples)
return img_pil
def instance_path(path, repo, file):
if repo == "":
if path == "none":
repo = "none"
else:
model_path = get_local_path(file_path, path, file)
repo = get_instance_path(model_path)
return repo
class MSdiffusion_Model_Loader:
def __init__(self):
self.counters = {}
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"clip_vision": ("CLIP_VISION",),
"ckpt_name": (["none"] + folder_paths.get_filename_list("checkpoints"),),
"vae_id": (["none"] + folder_paths.get_filename_list("vae"),),
"controlnet_ckpt": (["none"] + folder_paths.get_filename_list("controlnet"),),
"lora": (["none"] + folder_paths.get_filename_list("loras"),),
"lora_scale": ("FLOAT", {"default": 0.8, "min": 0.1, "max": 1.0, "step": 0.1}),
"trigger_words": ("STRING", {"default": "best quality"}),
"scheduler": (scheduler_list,), }
}
RETURN_TYPES = ("MSDIF_DICT",)
RETURN_NAMES = ("model",)
FUNCTION = "ms_model_loader"
CATEGORY = "MSdiffusion"
def ms_model_loader(self, clip_vision,ckpt_name, vae_id, controlnet_ckpt,
lora, lora_scale, trigger_words, scheduler):
image_encoder = clip_vision
device = image_encoder.load_device
scheduler_choice = get_scheduler(scheduler)
controlnet_path = folder_paths.get_full_path("controlnet", controlnet_ckpt)
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
modle_config = os.path.join(dir_path,"local_repo")
original_config_file = os.path.join(dir_path, "config", "sd_xl_base.yaml")
if lora != "none":
lora_path = folder_paths.get_full_path("loras", lora)
lora_path = get_instance_path(lora_path)
else:
lora_path = ""
if "/" in lora:
lora = lora.split("/")[-1]
if "\\" in lora:
lora = lora.split("\\")[-1]
# load SDXL pipeline
try:
pipe = StableDiffusionXLPipeline.from_single_file(
ckpt_path, config=modle_config, original_config=original_config_file, torch_dtype=torch.float16)
except:
try:
pipe = StableDiffusionXLPipeline.from_single_file(
ckpt_path, config=modle_config, original_config_file=original_config_file,
torch_dtype=torch.float16)
except:
raise "model error"
if controlnet_ckpt != "none": # 启用control
from safetensors.torch import load_file
controlnet = ControlNetModel.from_unet(pipe.unet)
cn_state_dict = load_file(controlnet_path, device="cpu")
controlnet.load_state_dict(cn_state_dict, strict=False)
controlnet.to(torch.float16)
pipe=StableDiffusionXLControlNetPipeline.from_pipe(pipe,controlnet=controlnet)
if vae_id != "none":
vae_id = folder_paths.get_full_path("vae", vae_id)
vae_config=os.path.join(modle_config,"vae")
pipe.vae = AutoencoderKL.from_single_file(vae_id, config=vae_config,torch_dtype=torch.float16).to(device)
if lora != "none":
if lora in lora_lightning_list:
pipe.load_lora_weights(lora_path)
pipe.fuse_lora()
else:
pipe.load_lora_weights(lora_path, adapter_name=trigger_words)
pipe.fuse_lora(adapter_names=[trigger_words, ], lora_scale=lora_scale)
pipe.scheduler = scheduler_choice.from_config(pipe.scheduler.config)
pipe.enable_xformers_memory_efficient_attention()
pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2)
pipe.enable_vae_slicing()
# if device != "mps":
# pipe.enable_model_cpu_offload()
ms_dir = os.path.join(file_path, "models", "photomaker")
photomaker_local_path = os.path.join(ms_dir, "ms_adapter.bin")
if not os.path.exists(photomaker_local_path):
ms_path = hf_hub_download(
repo_id="doge1516/MS-Diffusion",
filename="ms_adapter.bin",
repo_type="model",
local_dir=ms_dir,
)
else:
ms_path = photomaker_local_path
ms_ckpt = get_instance_path(ms_path)
image_processor = CLIPImageProcessor()
config_path = os.path.join(dir_path, "config", "config.json")
image_encoder_config = OmegaConf.load(config_path)
image_encoder_projection_dim = image_encoder_config["vision_config"]["projection_dim"]
num_tokens = 16
image_encoder_type = "clip"
image_proj_type = "resampler"
latent_init_mode = "grounding"
# latent_init_mode = "random"
image_proj_model = Resampler(
dim=1280,
depth=4,
dim_head=64,
heads=20,
num_queries=num_tokens,
embedding_dim=image_encoder_config["vision_config"]["hidden_size"], # image_encoder.config.hidden_size,
output_dim=pipe.unet.config.cross_attention_dim,
ff_mult=4,
latent_init_mode=latent_init_mode,
phrase_embeddings_dim=pipe.text_encoder.config.projection_dim,
).to(device, dtype=torch.float16)
ms_model = MSAdapter(pipe.unet, image_proj_model, ckpt_path=ms_ckpt, device=device, num_tokens=num_tokens)
ms_model.to(device, dtype=torch.float16)
pipe.to(device, dtype=torch.float16)
torch.cuda.empty_cache()
model = {"pipe": pipe, "ms_model": ms_model, "image_encoder": image_encoder, "image_processor": image_processor,
"lora": lora,
"trigger_words": trigger_words, "controlnet_model": controlnet_ckpt,
"image_encoder_type": image_encoder_type, "image_proj_type": image_proj_type,}
return (model,)
class MSdiffusion_Sampler:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"model": ("MSDIF_DICT",),
"prompt": ("STRING", {"multiline": True, "default": "a [dog] wearing a pink sunglass"}),
"negative_prompt": ("STRING", {"multiline": True,
"default": "(worst quality, low quality, normal quality, lowres),"}),
"seed": ("INT", {"default": 0, "min": 0, "max": MAX_SEED}),
"steps": ("INT", {"default": 20, "min": 1, "max": 100}),
"cfg": ("FLOAT", {"default": 7, "min": 0.1, "max": 10.0, "step": 0.1, "round": 0.01}),
"scale": (
"FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.1, "round": 0.01}),
"mask_threshold": (
"FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.1, "round": 0.01}),
"start_step": ("INT", {"default": 5, "min": 1, "max": 1024}),
"controlnet_scale": (
"FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.1, "round": 0.01}),
"width": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 64, "display": "number"}),
"height": ("INT", {"default": 768, "min": 256, "max": 2048, "step": 64, "display": "number"}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1, }),
"drop_grounding_tokens": ("BOOLEAN", {"default": False},),
"guidance_list": ("STRING", {"multiline": True, "default": "0., 0.25, 0.4, 0.75;0.6, 0.25, 1., 0.75"}),
},
"optional": {"control_image": ("IMAGE",),
}
}
@classmethod
def IS_CHANGED(cls, **kwargs):
if kwargs.get("batch_size"):
return
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("image",)
FUNCTION = "ms_sampler"
CATEGORY = "MSdiffusion"
def main_normal(self, prompt, pipe, phrases, ms_model, input_images, num_samples, steps, seed, negative_prompt,
scale, image_encoder, cfg, image_processor,
boxes, mask_threshold, start_step, image_proj_type, image_encoder_type, drop_grounding_tokens,
height, width, phrase_idxes, eot_idxes, image, ):
images = ms_model.generate(pipe=pipe, pil_images=[input_images], processed_images=image,
num_samples=num_samples,
num_inference_steps=steps,
seed=seed,
prompt=[prompt], negative_prompt=negative_prompt, scale=scale,
image_encoder=image_encoder, guidance_scale=cfg,
image_processor=image_processor, boxes=boxes,
mask_threshold=mask_threshold,
start_step=start_step,
image_proj_type=image_proj_type,
image_encoder_type=image_encoder_type,
phrases=phrases,
drop_grounding_tokens=drop_grounding_tokens,
phrase_idxes=phrase_idxes, eot_idxes=eot_idxes, height=height,
width=width)
return images
def main_control(self, prompt, width, height, pipe, phrases, ms_model, input_images, num_samples, steps, seed,
negative_prompt, scale, image_encoder, cfg,
image_processor, boxes, mask_threshold, start_step, image_proj_type, image_encoder_type,
drop_grounding_tokens, controlnet_scale, control_image, phrase_idxes, eot_idxes, image, ):
images = ms_model.generate(pipe=pipe, pil_images=[input_images], processed_images=image,
num_samples=num_samples,
num_inference_steps=steps,
seed=seed,
prompt=[prompt], negative_prompt=negative_prompt, scale=scale,
image_encoder=image_encoder, guidance_scale=cfg,
image_processor=image_processor, boxes=boxes,
mask_threshold=mask_threshold,
start_step=start_step,
image_proj_type=image_proj_type,
image_encoder_type=image_encoder_type,
phrases=phrases,
drop_grounding_tokens=drop_grounding_tokens,
phrase_idxes=phrase_idxes, eot_idxes=eot_idxes, height=height,
width=width,
image=control_image, controlnet_conditioning_scale=controlnet_scale)
return images
def get_phrases_idx(self, tokenizer, phrases, prompt):
res = []
phrase_cnt = {}
for phrase in phrases:
if phrase in phrase_cnt:
cur_cnt = phrase_cnt[phrase]
phrase_cnt[phrase] += 1
else:
cur_cnt = 0
phrase_cnt[phrase] = 1
res.append(get_phrase_idx(tokenizer, phrase, prompt, num=cur_cnt)[0])
return res
def get_float(self, str_in):
list_str = str_in.split(",")
float_box = [float(x) for x in list_str]
return float_box
def center_crop_s(self, img, new_width, new_height):
width, height = img.size
left = (width - new_width) / 2
top = (height - new_height) / 2
right = (width + new_width) / 2
bottom = (height + new_height) / 2
return img.crop((left, top, right, bottom))
def ms_sampler(self, image, model, prompt, negative_prompt, seed, steps,
cfg, scale, mask_threshold, start_step, controlnet_scale, width, height, batch_size,
drop_grounding_tokens, guidance_list, **kwargs):
if width==height:
new_width = width
new_height = width
min_wh= width
else:
min_wh=min(width,height)
max_wh=max(width,height)
new_width = max_wh
new_height =max_wh
pipe = model["pipe"]
ms_model = model["ms_model"]
image_encoder = model["image_encoder"]
image_processor = model["image_processor"]
lora = model["lora"]
trigger_words = model["trigger_words"]
controlnet_model = model["controlnet_model"]
image_encoder_type = model["image_encoder_type"]
image_proj_type = model["image_proj_type"]
if drop_grounding_tokens:
drop_grounding_tokens = [1] # set to 1 if you want to drop the grounding tokens
else:
drop_grounding_tokens = [0]
guidance_list = guidance_list.strip().split(";")
box_add = [] # 获取预设box
for i in range(len(guidance_list)):
box_add.append(self.get_float(guidance_list[i]))
if mask_threshold == 0.:
mask_threshold = None
if '[' in prompt and ']' in prompt:
object_prompt = extract_content_from_brackets(prompt) # 提取prompt的object list
object_prompt=[i.strip() for i in object_prompt]
for i in object_prompt:
if " " in i:
raise "when using [object],object must be a word,any blank in ititit will cause error."
object_prompt=[i for i in object_prompt ]
unique_object_prompt = sorted(list(set(object_prompt)),
key=lambda x: list(object_prompt).index(x)) # 清除同名物体,保持原有顺序
object_num = len(unique_object_prompt)
else:
raise "Need to enclose the name of the object in square brackets"
d1, _, _, _ = image.size()
# print(d1,object_num)
if lora != "none":
prompt = prompt + trigger_words + " " + "style"
prompt = prompt.replace("]", " ").replace("[", " ") # 去除[]
if object_num == d1:
phrases = [unique_object_prompt]
print(phrases)
else:
raise "The number of objects should be equal to the number of input images"
if d1 == 1:
input_images = [nomarl_upscale_topil(image, min_wh, min_wh)]
torch.cuda.empty_cache()
if mask_threshold: # step之前所用布局指导
boxes = box_add[0]
boxes = [[boxes]] # 随机选1个自定义的
# boxes = [[[0.25, 0.25, 0.75, 0.75]]] # 1 object
else:
boxes = [[[0., 0., 0., 0.]]] # 1 object
phrase_idxes = [self.get_phrases_idx(pipe.tokenizer, phrases[0], prompt)]
eot_idxes = [[get_eot_idx(pipe.tokenizer, prompt)] * len(phrases[0])]
if controlnet_model != "none":
control_image = kwargs.get("control_image")
control_img = nomarl_upscale_topil(control_image, min_wh, min_wh)
image_main = self.main_control(prompt, new_width, new_height, pipe, phrases, ms_model, input_images, batch_size,
steps, seed, negative_prompt, scale, image_encoder, cfg,
image_processor, boxes, mask_threshold, start_step, image_proj_type,
image_encoder_type, drop_grounding_tokens, controlnet_scale, control_img,
phrase_idxes, eot_idxes, image)
else: # no controlnet
image_main = self.main_normal(prompt, pipe, phrases, ms_model, input_images, batch_size, steps, seed,
negative_prompt, scale, image_encoder, cfg, image_processor,
boxes, mask_threshold, start_step, image_proj_type, image_encoder_type,
drop_grounding_tokens, new_height, new_width, phrase_idxes, eot_idxes, image
)
else:
img_list = list(torch.chunk(image, chunks=d1))
# print(img_list)
input_images = [nomarl_upscale_topil(img, min_wh, min_wh) for img in img_list]
# print(input_images)
if mask_threshold:
if len(box_add) < d1:
raise "When using 'mask_threshold', a sequence equal to the number of objects must be entered in 'guidance list'"
else:
boxes = [box_add[:d1]]
print(boxes)
# boxes = [[[0., 0.25, 0.4, 0.75], [0.6, 0.25, 1., 0.75]]] # man+women
else:
zero_list = [0 for _ in range(4)]
boxes = [zero_list for _ in range(d1)]
boxes = [boxes] # used if you want no layout guidance
# print(boxes)
phrase_idxes = [self.get_phrases_idx(pipe.tokenizer, phrases[0], prompt)]
eot_idxes = [[get_eot_idx(pipe.tokenizer, prompt)] * len(phrases[0])]
if controlnet_model != "none":
control_image = kwargs.get("control_image")
control_img = nomarl_upscale_topil(control_image, min_wh, min_wh) # pil
image_main = self.main_control(prompt, new_width, new_height, pipe, phrases, ms_model, input_images, batch_size,
steps, seed, negative_prompt, scale, image_encoder, cfg,
image_processor, boxes, mask_threshold, start_step, image_proj_type,
image_encoder_type, drop_grounding_tokens, controlnet_scale, control_img,
phrase_idxes, eot_idxes, image)
else: # no controlnet
image_main = self.main_normal(prompt, pipe, phrases, ms_model, input_images, batch_size, steps, seed,
negative_prompt, scale, image_encoder, cfg, image_processor,
boxes, mask_threshold, start_step, image_proj_type, image_encoder_type,
drop_grounding_tokens, new_height, new_width, phrase_idxes, eot_idxes, image
)
output_img = []
for i in range(batch_size):
if width!=height:
output_img.append(self.center_crop_s(image_main[i], width, height))
else:
output_img.append(image_main[i])
image_list = narry_list(output_img)
image = torch.from_numpy(np.fromiter(image_list, np.dtype((np.float32, (height, width, 3)))))
return (image,)
class MS_Object_img_Batch:
@classmethod
def INPUT_TYPES(s):
return {"required": {"image_a": ("IMAGE",),
},
"optional": {"image_b": ("IMAGE",),
"image_c": ("IMAGE",),
"image_d": ("IMAGE",)}
}
RETURN_TYPES = ("IMAGE",)
ETURN_NAMES = ("image",)
FUNCTION = "main_batch"
CATEGORY = "MSdiffusion"
def main_batch(self, image_a, **kwargs):
image_b = kwargs.get("image_b")
image_c = kwargs.get("image_c")
image_d = kwargs.get("image_d")
width = image_a.shape[2]
height = image_a.shape[1]
img_list = [image_a]
if isinstance(image_b, torch.Tensor):
image_b = nomarl_upscale(image_b, width, height)
img_list.append(image_b)
if isinstance(image_c, torch.Tensor):
image_c = nomarl_upscale(image_c, width, height)
img_list.append(image_c)
if isinstance(image_d, torch.Tensor):
image_d = nomarl_upscale(image_d, width, height)
img_list.append(image_d)
image = torch.cat(tuple(img_list), dim=0)
return (image,)
NODE_CLASS_MAPPINGS = {
"MSdiffusion_Model_Loader": MSdiffusion_Model_Loader,
"MSdiffusion_Sampler": MSdiffusion_Sampler,
"MS_Object_img_Batch": MS_Object_img_Batch,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MSdiffusion_Model_Loader": "MSdiffusion_Model_Loader",
"MSdiffusion_Sampler": "MSdiffusion_Sampler",
"MS_Object_img_Batch": "MS_Object_img_Batch"
}
README.md
+34 -73
View File
@@ -1,87 +1,48 @@
Youu can using MS_Diffusion in ComfyUI
# ComfyUI_MS_Diffusion
--
MS-Diffusion origin From: [MS-Diffusion](https://github.com/MS-Diffusion/MS-Diffusion)
----
[MS-Diffusion](https://github.com/MS-Diffusion/MS-Diffusion) multi-subject zero-shot image personalization with layout guidance
**NEW Update**
# Upadtes
* 借助agent（自动）使用comfyUI核心重新封装此方法 / Re-implement this method with the help of agent to re-package it into the core of comfyUI
---
* 2024/09/06:fix runway error/load single controlnet now /首次加载不再需要连外网。
**Previous updates**
* del clip repo，Add comfyUI clip_vision loader/加入comfyUI的clip vision节点，不再使用 clip repo。
1.Installation
# Installation
-----
In the ./ComfyUI /custom_node directory, run the following:
``` python
git clone https://github.com/smthemex/ComfyUI_MS_Diffusion.git
```
2.requirements
----
In the ./ComfyUI/custom_nodes directory, run the following:
```
pip install -r requirements.txt
```
缺啥装啥。
If the module is missing, please pip install
3 Need model
----
Need download "ms_adapter.bin" : [link](https://huggingface.co/doge1516/MS-Diffusion/tree/main)
clip vision model （ any base from CLIP-ViT-bigG-14-laion2B-39B-b160k）
Control_img image preprocessing, please use other nodes
```
├── ComfyUI/models/checkpoints/
| ├── any SDXL weigth
├── ComfyUI/models/photomaker/
| ├── ms_adapter.bin
├── ComfyUI/models/clip_vision/
| ├── clip_vision_g.safetensors or CLIP-ViT-bigG-14-laion2B-39B-b160k.safetensors
├── ComfyUI/models/controlnet
| ├──any SDXL controlnet
git clone https://github.com/smthemex/ComfyUI_MS_Diffusion
```
4 using tips
---
--生成物体名称需要用[ ]括起来，有多少个物体，就要有多少张图片输入;
--To generate object names, they need to be enclosed in [ ]. As many objects as there are, there must be as many images to input;
5 Example
----
normal 2 boject 常规双主体图生图 最新示例。
![](https://github.com/smthemex/ComfyUI_MS_Diffusion/blob/main/examples/example_new.png)
normal 4 boject 4主体图生图。更多的主体没有测试，旧示例，仅供参考。
![](https://github.com/smthemex/ComfyUI_MS_Diffusion/blob/main/examples/autolayerimg4img.png)
2 object zero shot and controlnet img2img 双物体加controlnet引导 图生图,旧示例，仅供参考
![](https://github.com/smthemex/ComfyUI_MS_Diffusion/blob/main/examples/controlnet%20_obj.png)
![](https://github.com/smthemex/ComfyUI_MS_Diffusion/blob/main/examples/controlnet%20_obj1.png)
6 Citation
------
MS-Diffusion
## Models
[sdxl](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
[ms-ipadapter](https://huggingface.co/doge1516/MS-Diffusion)
```
@misc{wang2024msdiffusion,
title={MS-Diffusion: Multi-subject Zero-shot Image Personalization with Layout Guidance},
author={X. Wang and Siming Fu and Qihan Huang and Wanggui He and Hao Jiang},
year={2024},
eprint={2406.07209},
archivePrefix={arXiv},
primaryClass={cs.CV}
}
├── ComfyUI/models/checkpoints
| ├── sd_xl_base_1.0.safetensors #any sdxl
├── ComfyUI/models/photomaker #ipadapter
| ├──ms_adapter.bin
```
IP-Adapter
## Multi-subject (multiple reference images)
| # | Sequence | Example (2 subjects) |
|---|---|---|
| 1 | IMAGE batch order (native `ImageBatch`: image1, image2, ...) | `ms_dog.jpg` → image1, `ms_cat.jpg` → image2 |
| 2 | `[subject]` order in `prompt` | `a [dog] and a [cat] on the beach` |
| 3 | `boxes`, `;`-separated, one `x1,y1,x2,y2` each | `0., 0.25, 0.4, 0.75;0.6, 0.25, 1., 0.75` |
## Example
![](https://github.com/smthemex/ComfyUI_MS_Diffusion/blob/main/example_workflows/example.png)
## Citation
```
@article{ye2023ip-adapter,
title={IP-Adapter: Text Compatible Image Prompt Adapter for Text-to-Image Diffusion Models},
author={Ye, Hu and Zhang, Jun and Liu, Sibo and Han, Xiao and Yang, Wei},
booktitle={arXiv preprint arxiv:2308.06721},
year={2023}
@inproceedings{
wang2025msdiffusion,
title={{MS}-Diffusion: Multi-subject Zero-shot Image Personalization with Layout Guidance},
author={Xierui Wang and Siming Fu and Qihan Huang and Wanggui He and Hao Jiang},
booktitle={The Thirteenth International Conference on Learning Representations},
year={2025},
url={https://openreview.net/forum?id=PJqP0wyQek}
}
```
__init__.py
+2 -2
View File
@@ -1,4 +1,4 @@
from .MS_Diffusion import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
from .msdiffusion_native import NODE_CLASS_MAPPINGS,NODE_DISPLAY_NAME_MAPPINGS
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
config/config.json
-184
View File
@@ -1,184 +0,0 @@
{
"_commit_hash": null,
"architectures": [
"CLIPModel"
],
"initializer_factor": 1.0,
"logit_scale_init_value": 2.6592,
"model_type": "clip",
"projection_dim": 1280,
"text_config": {
"_name_or_path": "",
"add_cross_attention": false,
"architectures": null,
"attention_dropout": 0.0,
"bad_words_ids": null,
"begin_suppress_tokens": null,
"bos_token_id": 0,
"chunk_size_feed_forward": 0,
"cross_attention_hidden_size": null,
"decoder_start_token_id": null,
"diversity_penalty": 0.0,
"do_sample": false,
"dropout": 0.0,
"early_stopping": false,
"encoder_no_repeat_ngram_size": 0,
"eos_token_id": 2,
"exponential_decay_length_penalty": null,
"finetuning_task": null,
"forced_bos_token_id": null,
"forced_eos_token_id": null,
"hidden_act": "gelu",
"hidden_size": 1280,
"id2label": {
"0": "LABEL_0",
"1": "LABEL_1"
},
"initializer_factor": 1.0,
"initializer_range": 0.02,
"intermediate_size": 5120,
"is_decoder": false,
"is_encoder_decoder": false,
"label2id": {
"LABEL_0": 0,
"LABEL_1": 1
},
"layer_norm_eps": 1e-05,
"length_penalty": 1.0,
"max_length": 20,
"max_position_embeddings": 77,
"min_length": 0,
"model_type": "clip_text_model",
"no_repeat_ngram_size": 0,
"num_attention_heads": 20,
"num_beam_groups": 1,
"num_beams": 1,
"num_hidden_layers": 32,
"num_return_sequences": 1,
"output_attentions": false,
"output_hidden_states": false,
"output_scores": false,
"pad_token_id": 1,
"prefix": null,
"problem_type": null,
"projection_dim": 1280,
"pruned_heads": {},
"remove_invalid_values": false,
"repetition_penalty": 1.0,
"return_dict": true,
"return_dict_in_generate": false,
"sep_token_id": null,
"suppress_tokens": null,
"task_specific_params": null,
"temperature": 1.0,
"tf_legacy_loss": false,
"tie_encoder_decoder": false,
"tie_word_embeddings": true,
"tokenizer_class": null,
"top_k": 50,
"top_p": 1.0,
"torch_dtype": null,
"torchscript": false,
"transformers_version": "4.24.0",
"typical_p": 1.0,
"use_bfloat16": false,
"vocab_size": 49408
},
"text_config_dict": {
"hidden_act": "gelu",
"hidden_size": 1280,
"intermediate_size": 5120,
"num_attention_heads": 20,
"num_hidden_layers": 32,
"projection_dim": 1280
},
"torch_dtype": "float32",
"transformers_version": null,
"vision_config": {
"_name_or_path": "",
"add_cross_attention": false,
"architectures": null,
"attention_dropout": 0.0,
"bad_words_ids": null,
"begin_suppress_tokens": null,
"bos_token_id": null,
"chunk_size_feed_forward": 0,
"cross_attention_hidden_size": null,
"decoder_start_token_id": null,
"diversity_penalty": 0.0,
"do_sample": false,
"dropout": 0.0,
"early_stopping": false,
"encoder_no_repeat_ngram_size": 0,
"eos_token_id": null,
"exponential_decay_length_penalty": null,
"finetuning_task": null,
"forced_bos_token_id": null,
"forced_eos_token_id": null,
"hidden_act": "gelu",
"hidden_size": 1664,
"id2label": {
"0": "LABEL_0",
"1": "LABEL_1"
},
"image_size": 224,
"initializer_factor": 1.0,
"initializer_range": 0.02,
"intermediate_size": 8192,
"is_decoder": false,
"is_encoder_decoder": false,
"label2id": {
"LABEL_0": 0,
"LABEL_1": 1
},
"layer_norm_eps": 1e-05,
"length_penalty": 1.0,
"max_length": 20,
"min_length": 0,
"model_type": "clip_vision_model",
"no_repeat_ngram_size": 0,
"num_attention_heads": 16,
"num_beam_groups": 1,
"num_beams": 1,
"num_channels": 3,
"num_hidden_layers": 48,
"num_return_sequences": 1,
"output_attentions": false,
"output_hidden_states": false,
"output_scores": false,
"pad_token_id": null,
"patch_size": 14,
"prefix": null,
"problem_type": null,
"projection_dim": 1280,
"pruned_heads": {},
"remove_invalid_values": false,
"repetition_penalty": 1.0,
"return_dict": true,
"return_dict_in_generate": false,
"sep_token_id": null,
"suppress_tokens": null,
"task_specific_params": null,
"temperature": 1.0,
"tf_legacy_loss": false,
"tie_encoder_decoder": false,
"tie_word_embeddings": true,
"tokenizer_class": null,
"top_k": 50,
"top_p": 1.0,
"torch_dtype": null,
"torchscript": false,
"transformers_version": "4.24.0",
"typical_p": 1.0,
"use_bfloat16": false
},
"vision_config_dict": {
"hidden_act": "gelu",
"hidden_size": 1664,
"intermediate_size": 8192,
"num_attention_heads": 16,
"num_hidden_layers": 48,
"patch_size": 14,
"projection_dim": 1280
}
}
config/lora.yaml
-33
View File
@@ -1,33 +0,0 @@
lightning_unet:
- sdxl_lightning_1step_unet_x0.safetensors #repo ByteDance/SDXL-Lightning
- sdxl_lightning_2step_unet.safetensors #repo ByteDance/SDXL-Lightning
- sdxl_lightning_4step_unet.safetensors
- sdxl_lightning_8step_unet.safetensors
- Hyper-SDXL-1step-Unet.safetensors #repo ByteDance/Hyper-SD
- lcm-sdxl-base-1.0.safetensors #repo ckpt/lcm-sdxl-unet you can change model name like example
- dmd2_sdxl_1step_unet_fp16.bin #repo tianweiy/DMD2
- dmd2_sdxl_4step_unet_fp16.bin
lightning_xl_lora:
- Hyper-SDXL-12steps-CFG-lora.safetensors
- Hyper-SDXL-1step-lora.safetensors
- Hyper-SDXL-2step-lora.safetensors
- Hyper-SDXL-4step-lora.safetensors
- Hyper-SDXL-8step-lora.safetensors
- Hyper-SDXL-8steps-CFG-lora.safetensors
- sdxl_lightning_2step_lora.safetensors
- sdxl_lightning_4step_lora.safetensors
- sdxl_lightning_8step_lora.safetensors
- pcm_sdxl_lcmlike_lora_converted.safetensors
- pcm_sdxl_normalcfg_16step_converted.safetensors
- pcm_sdxl_normalcfg_4step_converted.safetensors
- pcm_sdxl_normalcfg_8step_converted.safetensors
- pcm_sdxl_smallcfg_16step_converted.safetensors
- pcm_sdxl_smallcfg_2step_converted.safetensors
- pcm_sdxl_smallcfg_4step_converted.safetensors
- pcm_sdxl_smallcfg_8step_converted.safetensors
- lcm-lora-sdxl.safetensors
- dmd2_sdxl_4step_lora.safetensors
- dmd2_sdxl_4step_lora_fp16.safetensors
- TCD-SDXL-LoRA.safetensors #need rename and TCD
- manne_turbo.safetensors
config/sd_xl_base.yaml
-93
View File
@@ -1,93 +0,0 @@
model:
target: sgm.models.diffusion.DiffusionEngine
params:
scale_factor: 0.13025
disable_first_stage_autocast: True
denoiser_config:
target: sgm.modules.diffusionmodules.denoiser.DiscreteDenoiser
params:
num_idx: 1000
scaling_config:
target: sgm.modules.diffusionmodules.denoiser_scaling.EpsScaling
discretization_config:
target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
network_config:
target: sgm.modules.diffusionmodules.openaimodel.UNetModel
params:
adm_in_channels: 2816
num_classes: sequential
use_checkpoint: True
in_channels: 4
out_channels: 4
model_channels: 320
attention_resolutions: [4, 2]
num_res_blocks: 2
channel_mult: [1, 2, 4]
num_head_channels: 64
use_linear_in_transformer: True
transformer_depth: [1, 2, 10]
context_dim: 2048
spatial_transformer_attn_type: softmax-xformers
conditioner_config:
target: sgm.modules.GeneralConditioner
params:
emb_models:
- is_trainable: False
input_key: txt
target: sgm.modules.encoders.modules.FrozenCLIPEmbedder
params:
layer: hidden
layer_idx: 11
- is_trainable: False
input_key: txt
target: sgm.modules.encoders.modules.FrozenOpenCLIPEmbedder2
params:
arch: ViT-bigG-14
version: laion2b_s39b_b160k
freeze: True
layer: penultimate
always_return_pooled: True
legacy: False
- is_trainable: False
input_key: original_size_as_tuple
target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
params:
outdim: 256
- is_trainable: False
input_key: crop_coords_top_left
target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
params:
outdim: 256
- is_trainable: False
input_key: target_size_as_tuple
target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
params:
outdim: 256
first_stage_config:
target: sgm.models.autoencoder.AutoencoderKL
params:
embed_dim: 4
monitor: val/rec_loss
ddconfig:
attn_type: vanilla-xformers
double_z: true
z_channels: 4
resolution: 256
in_channels: 3
out_ch: 3
ch: 128
ch_mult: [1, 2, 4, 4]
num_res_blocks: 2
attn_resolutions: []
dropout: 0.0
lossconfig:
target: torch.nn.Identity
example_workflows/example.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 472 KiB

example_workflows/msdiffusion.json
+864
View File
@@ -0,0 +1,864 @@
{
"id": "86c7c72a-3cc3-4d80-ac53-395f0f1a672f",
"revision": 0,
"last_node_id": 19,
"last_link_id": 44,
"nodes": [
{
"id": 2,
"type": "CLIPVisionLoader",
"pos": [
370.8325188848841,
179.37975031596505
],
"size": [
288.5261157862166,
61.99996640263096
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"label": "CLIP_VISION",
"name": "CLIP_VISION",
"type": "CLIP_VISION",
"slot_index": 0,
"links": [
1
]
}
],
"title": "CLIP Vision G",
"properties": {
"Node name for S&R": "CLIPVisionLoader"
},
"widgets_values": [
"clip_vision_g.safetensors"
],
"widgets_values_named": {
"clip_name": "clip_vision_g.safetensors"
}
},
{
"id": 3,
"type": "LoadImage",
"pos": [
370.7321179172254,
417.8313576180781
],
"size": [
315,
120
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"label": "IMAGE",
"name": "IMAGE",
"type": "IMAGE",
"slot_index": 0,
"links": [
20
]
},
{
"label": "MASK",
"name": "MASK",
"type": "MASK",
"slot_index": 1,
"links": null
}
],
"title": "subject image(s)",
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"rabbit.png",
"image"
],
"widgets_values_named": {
"image": "rabbit.png",
"upload": "image"
}
},
{
"id": 4,
"type": "CLIPVisionEncode",
"pos": [
388.2932244009372,
296.3283550695895
],
"size": [
297.9367589978424,
85.87358125758468
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"label": "clip_vision",
"name": "clip_vision",
"type": "CLIP_VISION",
"link": 1
},
{
"label": "image",
"name": "image",
"type": "IMAGE",
"link": 20
}
],
"outputs": [
{
"label": "CLIP_VISION_OUTPUT",
"name": "CLIP_VISION_OUTPUT",
"type": "CLIP_VISION_OUTPUT",
"slot_index": 0,
"links": [
37
]
}
],
"title": "CLIP Vision encode",
"properties": {
"Node name for S&R": "CLIPVisionEncode"
},
"widgets_values": [
"center"
],
"widgets_values_named": {
"crop": "center"
}
},
{
"id": 5,
"type": "MSDiffusionLoader",
"pos": [
366.4898045653227,
-103.33187302053278
],
"size": [
314.0874241798911,
85.99999774294434
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [
{
"label": "MS_ADAPTER",
"name": "ms_adapter",
"type": "MS_ADAPTER",
"slot_index": 0,
"links": [
30,
33
]
}
],
"title": "MS adapter",
"properties": {
"Node name for S&R": "MSDiffusionLoader"
},
"widgets_values": [
"ms_adapter.bin"
],
"widgets_values_named": {
"ms_adapter": "ms_adapter.bin"
},
"color": "#322",
"bgcolor": "#533"
},
{
"id": 7,
"type": "CLIPTextEncode",
"pos": [
789.8543229164497,
104.03036582957571
],
"size": [
261.1503022979092,
57.99999888788844
],
"flags": {
"collapsed": true
},
"order": 10,
"mode": 0,
"inputs": [
{
"label": "clip",
"name": "clip",
"type": "CLIP",
"link": 42
},
{
"name": "text",
"type": "STRING",
"widget": {
"name": "text"
},
"link": 44
}
],
"outputs": [
{
"label": "CONDITIONING",
"name": "CONDITIONING",
"type": "CONDITIONING",
"slot_index": 0,
"links": [
23
]
}
],
"title": "positive",
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"best quality, high quality, a dog on the beach"
],
"widgets_values_named": {
"text": "best quality, high quality, a dog on the beach"
}
},
{
"id": 8,
"type": "CLIPTextEncode",
"pos": [
789.7131145434341,
155.84224816498624
],
"size": [
290.3512025062168,
101.9999896298381
],
"flags": {
"collapsed": true
},
"order": 6,
"mode": 0,
"inputs": [
{
"label": "clip",
"name": "clip",
"type": "CLIP",
"link": 43
}
],
"outputs": [
{
"label": "CONDITIONING",
"name": "CONDITIONING",
"type": "CONDITIONING",
"slot_index": 0,
"links": [
21
]
}
],
"title": "negative",
"properties": {
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"(worst quality, low quality, normal quality, lowres),"
],
"widgets_values_named": {
"text": "(worst quality, low quality, normal quality, lowres),"
}
},
{
"id": 10,
"type": "EmptyLatentImage",
"pos": [
743.7407173134332,
533.9610772211638
],
"size": [
341.4616713091804,
127.29972476947185
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [],
"outputs": [
{
"label": "LATENT",
"name": "LATENT",
"type": "LATENT",
"slot_index": 0,
"links": [
28
]
}
],
"title": "empty latent",
"properties": {
"Node name for S&R": "EmptyLatentImage"
},
"widgets_values": [
1024,
1024,
1
],
"widgets_values_named": {
"width": 1024,
"height": 1024,
"batch_size": 1
}
},
{
"id": 11,
"type": "KSampler",
"pos": [
1156.9620824336098,
80.54740592847531
],
"size": [
315,
262
],
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"label": "model",
"name": "model",
"type": "MODEL",
"link": 31
},
{
"label": "positive",
"name": "positive",
"type": "CONDITIONING",
"link": 23
},
{
"label": "negative",
"name": "negative",
"type": "CONDITIONING",
"link": 21
},
{
"label": "latent_image",
"name": "latent_image",
"type": "LATENT",
"link": 28
}
],
"outputs": [
{
"label": "LATENT",
"name": "LATENT",
"type": "LATENT",
"slot_index": 0,
"links": [
15
]
}
],
"title": "KSampler",
"properties": {
"Node name for S&R": "KSampler"
},
"widgets_values": [
0,
"fixed",
20,
7,
"euler",
"normal",
1
],
"widgets_values_named": {
"seed": 0,
"control_after_generate": "fixed",
"steps": 20,
"cfg": 7,
"sampler_name": "euler",
"scheduler": "normal",
"denoise": 1
}
},
{
"id": 12,
"type": "VAEDecode",
"pos": [
1161.771010927863,
430.2061194051902
],
"size": [
235.70637510865785,
54.000003526763365
],
"flags": {},
"order": 12,
"mode": 0,
"inputs": [
{
"label": "samples",
"name": "samples",
"type": "LATENT",
"link": 15
},
{
"label": "vae",
"name": "vae",
"type": "VAE",
"link": 41
}
],
"outputs": [
{
"label": "IMAGE",
"name": "IMAGE",
"type": "IMAGE",
"slot_index": 0,
"links": [
17
]
}
],
"title": "VAE decode",
"properties": {
"Node name for S&R": "VAEDecode"
}
},
{
"id": 13,
"type": "SaveImage",
"pos": [
1436.9519060913603,
380.93267297794523
],
"size": [
323.21237046819726,
345.78613186594066
],
"flags": {},
"order": 13,
"mode": 0,
"inputs": [
{
"label": "images",
"name": "images",
"type": "IMAGE",
"link": 17
}
],
"outputs": [
{
"name": "images",
"type": "IMAGE",
"links": null
}
],
"title": "save",
"properties": {
"Node name for S&R": "SaveImage"
},
"widgets_values": [
"MSDiffusion_native"
],
"widgets_values_named": {
"filename_prefix": "MSDiffusion_native"
}
},
{
"id": 15,
"type": "MSDiffusionApply",
"pos": [
1162.9935239715307,
-124.92663717491732
],
"size": [
312.0376519607282,
153.9999850172292
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 39
},
{
"name": "ms_adapter",
"type": "MS_ADAPTER",
"link": 33
},
{
"name": "ms_embeds",
"type": "MS_EMBEDS",
"link": 35
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
31
]
}
],
"properties": {
"Node name for S&R": "MSDiffusionApply"
},
"widgets_values": [
0.8,
5
],
"widgets_values_named": {
"scale": 0.8,
"start_step": 5
},
"color": "#322",
"bgcolor": "#533"
},
{
"id": 16,
"type": "MSDiffusionConditioning",
"pos": [
730.952378177398,
206.97648380148877
],
"size": [
368.06279843507633,
290.0000117777662
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "ms_adapter",
"type": "MS_ADAPTER",
"link": 30
},
{
"name": "clip",
"type": "CLIP",
"link": 40
},
{
"name": "clip_vision_output",
"type": "CLIP_VISION_OUTPUT",
"link": 37
}
],
"outputs": [
{
"name": "ms_embeds",
"type": "MS_EMBEDS",
"links": [
35
]
},
{
"name": "prompt_clean",
"type": "STRING",
"links": [
36,
44
]
}
],
"properties": {
"Node name for S&R": "MSDiffusionConditioning"
},
"widgets_values": [
"best quality, high quality, a [rabbit] on the beach",
"0.25, 0.25, 0.75, 0.75",
0.5,
false
],
"widgets_values_named": {
"prompt": "best quality, high quality, a [rabbit] on the beach",
"boxes": "0.25, 0.25, 0.75, 0.75",
"mask_threshold": 0.5,
"drop_grounding_tokens": false
},
"color": "#322",
"bgcolor": "#533"
},
{
"id": 18,
"type": "PreviewAny",
"pos": [
1154.0802681233788,
536.8622571930127
],
"size": [
250.7532473108919,
138.00003538379124
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "source",
"type": "*",
"link": 36
}
],
"outputs": [
{
"name": "STRING",
"type": "STRING",
"links": null
}
],
"properties": {
"Node name for S&R": "PreviewAny"
},
"widgets_values": [],
"widgets_values_named": {}
},
{
"id": 19,
"type": "CheckpointLoaderSimple",
"pos": [
363.48933959561873,
32.07108167105924
],
"size": [
303.36346898987495,
102.00001573582173
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
39
]
},
{
"name": "CLIP",
"type": "CLIP",
"links": [
40,
42,
43
]
},
{
"name": "VAE",
"type": "VAE",
"links": [
41
]
}
],
"properties": {
"Node name for S&R": "CheckpointLoaderSimple"
},
"widgets_values": [
"sd_xl_base_1.0.safetensors"
],
"widgets_values_named": {
"ckpt_name": "sd_xl_base_1.0.safetensors"
}
}
],
"links": [
[
1,
2,
0,
4,
0,
"CLIP_VISION"
],
[
15,
11,
0,
12,
0,
"LATENT"
],
[
17,
12,
0,
13,
0,
"IMAGE"
],
[
20,
3,
0,
4,
1,
"IMAGE"
],
[
21,
8,
0,
11,
2,
"CONDITIONING"
],
[
23,
7,
0,
11,
1,
"CONDITIONING"
],
[
28,
10,
0,
11,
3,
"LATENT"
],
[
30,
5,
0,
16,
0,
"MS_ADAPTER"
],
[
31,
15,
0,
11,
0,
"MODEL"
],
[
33,
5,
0,
15,
1,
"MS_ADAPTER"
],
[
35,
16,
0,
15,
2,
"MS_EMBEDS"
],
[
36,
16,
1,
18,
0,
"STRING"
],
[
37,
4,
0,
16,
2,
"CLIP_VISION_OUTPUT"
],
[
39,
19,
0,
15,
0,
"MODEL"
],
[
40,
19,
1,
16,
1,
"CLIP"
],
[
41,
19,
2,
12,
1,
"VAE"
],
[
42,
19,
1,
7,
0,
"CLIP"
],
[
43,
19,
1,
8,
0,
"CLIP"
],
[
44,
16,
1,
7,
1,
"STRING"
]
],
"groups": [
{
"id": 1,
"title": "Group",
"bounding": [
353.48933959561873,
-187.6268820305459,
1416.6749369639388,
937.4582396486239
],
"color": "#3f789e",
"flags": {}
}
],
"config": {},
"extra": {
"ds": {
"scale": 0.964414990215265,
"offset": [
-61.41614699554228,
257.04284223878574
]
},
"frontendVersion": "1.58.5"
},
"version": 0.4
}
examples/autolayerimg2img.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 450 KiB

examples/autolayerimg4img.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 420 KiB

examples/controlnet _obj.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 460 KiB

examples/controlnet _obj1.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 469 KiB

examples/example_new.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 362 KiB

examples/workflow_new.json
-395
View File
@@ -1,395 +0,0 @@
{
"last_node_id": 40,
"last_link_id": 101,
"nodes": [
{
"id": 3,
"type": "LoadImage",
"pos": [
1119,
-33
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
89
],
"shape": 3,
"label": "IMAGE",
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3,
"label": "MASK"
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"123 (1).png",
"image"
]
},
{
"id": 36,
"type": "MS_Object_img_Batch",
"pos": [
1512,
180
],
"size": {
"0": 210,
"1": 86
},
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "image_a",
"type": "IMAGE",
"link": 88,
"label": "image_a"
},
{
"name": "image_b",
"type": "IMAGE",
"link": 89,
"label": "image_b"
},
{
"name": "image_c",
"type": "IMAGE",
"link": null,
"label": "image_c"
},
{
"name": "image_d",
"type": "IMAGE",
"link": null,
"label": "image_d"
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
98
],
"shape": 3,
"label": "IMAGE",
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "MS_Object_img_Batch"
}
},
{
"id": 4,
"type": "SaveImage",
"pos": [
1939,
362
],
"size": {
"0": 315,
"1": 270
},
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 99,
"label": "images"
}
],
"properties": {},
"widgets_values": [
"ComfyUI"
]
},
{
"id": 39,
"type": "MSdiffusion_Sampler",
"pos": [
1508,
349
],
"size": {
"0": 400,
"1": 504
},
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 98,
"label": "image"
},
{
"name": "model",
"type": "MSDIF_DICT",
"link": 101,
"label": "model"
},
{
"name": "control_image",
"type": "IMAGE",
"link": null,
"label": "control_image"
}
],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": [
99
],
"shape": 3,
"label": "image",
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "MSdiffusion_Sampler"
},
"widgets_values": [
"a [pig] wearing a pink [sunglass]",
"(worst quality, low quality, normal quality, lowres),",
825964470,
"fixed",
20,
7,
0.8,
0.5,
5,
0.8,
768,
768,
1,
false,
"0., 0.25, 0.4, 0.75;0., 0.25, 0.4, 0.75"
]
},
{
"id": 35,
"type": "LoadImage",
"pos": [
774,
-35
],
"size": {
"0": 315,
"1": 314
},
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
88
],
"shape": 3,
"label": "IMAGE",
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"shape": 3,
"label": "MASK"
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"123.png",
"image"
]
},
{
"id": 32,
"type": "CLIPVisionLoader",
"pos": [
790,
448
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 2,
"mode": 0,
"outputs": [
{
"name": "CLIP_VISION",
"type": "CLIP_VISION",
"links": [
100
],
"shape": 3,
"label": "CLIP_VISION",
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "CLIPVisionLoader"
},
"widgets_values": [
"clip_vision_g.safetensors"
]
},
{
"id": 40,
"type": "MSdiffusion_Model_Loader",
"pos": [
1162,
402
],
"size": {
"0": 315,
"1": 250
},
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "clip_vision",
"type": "CLIP_VISION",
"link": 100,
"label": "clip_vision"
}
],
"outputs": [
{
"name": "model",
"type": "MSDIF_DICT",
"links": [
101
],
"shape": 3,
"label": "model",
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "MSdiffusion_Model_Loader"
},
"widgets_values": [
"0SDXL\\juggernautXL_v8Rundiffusion.safetensors",
"",
"none",
"none",
"",
"none",
0.8,
"best quality",
"Euler"
]
}
],
"links": [
[
88,
35,
0,
36,
0,
"IMAGE"
],
[
89,
3,
0,
36,
1,
"IMAGE"
],
[
98,
36,
0,
39,
0,
"IMAGE"
],
[
99,
39,
0,
4,
0,
"IMAGE"
],
[
100,
32,
0,
40,
0,
"CLIP_VISION"
],
[
101,
40,
0,
39,
1,
"MSDIF_DICT"
]
],
"groups": [
{
"title": "Group",
"bounding": [
750,
-184,
1544,
1093
],
"color": "#3f789e",
"font_size": 24,
"locked": false
}
],
"config": {},
"extra": {
"ds": {
"scale": 0.9090909090909091,
"offset": {
"0": -600.3953810605899,
"1": 110.50656660086888
}
}
},
"version": 0.4
}
local_repo/model_index.json
-34
View File
@@ -1,34 +0,0 @@
{
"_class_name": "StableDiffusionXLPipeline",
"_diffusers_version": "0.19.0.dev0",
"force_zeros_for_empty_prompt": true,
"add_watermarker": null,
"scheduler": [
"diffusers",
"EulerDiscreteScheduler"
],
"text_encoder": [
"transformers",
"CLIPTextModel"
],
"text_encoder_2": [
"transformers",
"CLIPTextModelWithProjection"
],
"tokenizer": [
"transformers",
"CLIPTokenizer"
],
"tokenizer_2": [
"transformers",
"CLIPTokenizer"
],
"unet": [
"diffusers",
"UNet2DConditionModel"
],
"vae": [
"diffusers",
"AutoencoderKL"
]
}
local_repo/scheduler/scheduler_config.json
-18
View File
@@ -1,18 +0,0 @@
{
"_class_name": "EulerDiscreteScheduler",
"_diffusers_version": "0.19.0.dev0",
"beta_end": 0.012,
"beta_schedule": "scaled_linear",
"beta_start": 0.00085,
"clip_sample": false,
"interpolation_type": "linear",
"num_train_timesteps": 1000,
"prediction_type": "epsilon",
"sample_max_value": 1.0,
"set_alpha_to_one": false,
"skip_prk_steps": true,
"steps_offset": 1,
"timestep_spacing": "leading",
"trained_betas": null,
"use_karras_sigmas": false
}
local_repo/text_encoder/config.json
-24
View File
@@ -1,24 +0,0 @@
{
"architectures": [
"CLIPTextModel"
],
"attention_dropout": 0.0,
"bos_token_id": 0,
"dropout": 0.0,
"eos_token_id": 2,
"hidden_act": "quick_gelu",
"hidden_size": 768,
"initializer_factor": 1.0,
"initializer_range": 0.02,
"intermediate_size": 3072,
"layer_norm_eps": 1e-05,
"max_position_embeddings": 77,
"model_type": "clip_text_model",
"num_attention_heads": 12,
"num_hidden_layers": 12,
"pad_token_id": 1,
"projection_dim": 768,
"torch_dtype": "float16",
"transformers_version": "4.32.0.dev0",
"vocab_size": 49408
}
local_repo/text_encoder_2/config.json
-24
View File
@@ -1,24 +0,0 @@
{
"architectures": [
"CLIPTextModelWithProjection"
],
"attention_dropout": 0.0,
"bos_token_id": 0,
"dropout": 0.0,
"eos_token_id": 2,
"hidden_act": "gelu",
"hidden_size": 1280,
"initializer_factor": 1.0,
"initializer_range": 0.02,
"intermediate_size": 5120,
"layer_norm_eps": 1e-05,
"max_position_embeddings": 77,
"model_type": "clip_text_model",
"num_attention_heads": 20,
"num_hidden_layers": 32,
"pad_token_id": 1,
"projection_dim": 1280,
"torch_dtype": "float16",
"transformers_version": "4.32.0.dev0",
"vocab_size": 49408
}
local_repo/tokenizer/merges.txt
-48895
View File
File diff suppressed because it is too large Load Diff
local_repo/tokenizer/special_tokens_map.json
-24
View File
@@ -1,24 +0,0 @@
{
"bos_token": {
"content": "<|startoftext|>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
},
"eos_token": {
"content": "<|endoftext|>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
},
"pad_token": "<|endoftext|>",
"unk_token": {
"content": "<|endoftext|>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
}
}
local_repo/tokenizer/tokenizer_config.json
-33
View File
@@ -1,33 +0,0 @@
{
"add_prefix_space": false,
"bos_token": {
"__type": "AddedToken",
"content": "<|startoftext|>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
},
"clean_up_tokenization_spaces": true,
"do_lower_case": true,
"eos_token": {
"__type": "AddedToken",
"content": "<|endoftext|>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
},
"errors": "replace",
"model_max_length": 77,
"pad_token": "<|endoftext|>",
"tokenizer_class": "CLIPTokenizer",
"unk_token": {
"__type": "AddedToken",
"content": "<|endoftext|>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
}
}
local_repo/tokenizer/vocab.json
-49410
View File
File diff suppressed because it is too large Load Diff
local_repo/tokenizer_2/merges.txt
-48895
View File
File diff suppressed because it is too large Load Diff
local_repo/tokenizer_2/special_tokens_map.json
-24
View File
@@ -1,24 +0,0 @@
{
"bos_token": {
"content": "<|startoftext|>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
},
"eos_token": {
"content": "<|endoftext|>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
},
"pad_token": "!",
"unk_token": {
"content": "<|endoftext|>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
}
}
local_repo/tokenizer_2/tokenizer_config.json
-33
View File
@@ -1,33 +0,0 @@
{
"add_prefix_space": false,
"bos_token": {
"__type": "AddedToken",
"content": "<|startoftext|>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
},
"clean_up_tokenization_spaces": true,
"do_lower_case": true,
"eos_token": {
"__type": "AddedToken",
"content": "<|endoftext|>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
},
"errors": "replace",
"model_max_length": 77,
"pad_token": "!",
"tokenizer_class": "CLIPTokenizer",
"unk_token": {
"__type": "AddedToken",
"content": "<|endoftext|>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
}
}
local_repo/tokenizer_2/vocab.json
-49410
View File
File diff suppressed because it is too large Load Diff
local_repo/unet/config.json
-69
View File
@@ -1,69 +0,0 @@
{
"_class_name": "UNet2DConditionModel",
"_diffusers_version": "0.19.0.dev0",
"act_fn": "silu",
"addition_embed_type": "text_time",
"addition_embed_type_num_heads": 64,
"addition_time_embed_dim": 256,
"attention_head_dim": [
5,
10,
20
],
"block_out_channels": [
320,
640,
1280
],
"center_input_sample": false,
"class_embed_type": null,
"class_embeddings_concat": false,
"conv_in_kernel": 3,
"conv_out_kernel": 3,
"cross_attention_dim": 2048,
"cross_attention_norm": null,
"down_block_types": [
"DownBlock2D",
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D"
],
"downsample_padding": 1,
"dual_cross_attention": false,
"encoder_hid_dim": null,
"encoder_hid_dim_type": null,
"flip_sin_to_cos": true,
"freq_shift": 0,
"in_channels": 4,
"layers_per_block": 2,
"mid_block_only_cross_attention": null,
"mid_block_scale_factor": 1,
"mid_block_type": "UNetMidBlock2DCrossAttn",
"norm_eps": 1e-05,
"norm_num_groups": 32,
"num_attention_heads": null,
"num_class_embeds": null,
"only_cross_attention": false,
"out_channels": 4,
"projection_class_embeddings_input_dim": 2816,
"resnet_out_scale_factor": 1.0,
"resnet_skip_time_act": false,
"resnet_time_scale_shift": "default",
"sample_size": 128,
"time_cond_proj_dim": null,
"time_embedding_act_fn": null,
"time_embedding_dim": null,
"time_embedding_type": "positional",
"timestep_post_act": null,
"transformer_layers_per_block": [
1,
2,
10
],
"up_block_types": [
"CrossAttnUpBlock2D",
"CrossAttnUpBlock2D",
"UpBlock2D"
],
"upcast_attention": null,
"use_linear_projection": true
}
local_repo/vae/config.json
-32
View File
@@ -1,32 +0,0 @@
{
"_class_name": "AutoencoderKL",
"_diffusers_version": "0.20.0.dev0",
"_name_or_path": "../sdxl-vae/",
"act_fn": "silu",
"block_out_channels": [
128,
256,
512,
512
],
"down_block_types": [
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D"
],
"force_upcast": true,
"in_channels": 3,
"latent_channels": 4,
"layers_per_block": 2,
"norm_num_groups": 32,
"out_channels": 3,
"sample_size": 1024,
"scaling_factor": 0.13025,
"up_block_types": [
"UpDecoderBlock2D",
"UpDecoderBlock2D",
"UpDecoderBlock2D",
"UpDecoderBlock2D"
]
}
local_repo/vae_1_0/config.json
-31
View File
@@ -1,31 +0,0 @@
{
"_class_name": "AutoencoderKL",
"_diffusers_version": "0.19.0.dev0",
"act_fn": "silu",
"block_out_channels": [
128,
256,
512,
512
],
"down_block_types": [
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D"
],
"force_upcast": true,
"in_channels": 3,
"latent_channels": 4,
"layers_per_block": 2,
"norm_num_groups": 32,
"out_channels": 3,
"sample_size": 1024,
"scaling_factor": 0.13025,
"up_block_types": [
"UpDecoderBlock2D",
"UpDecoderBlock2D",
"UpDecoderBlock2D",
"UpDecoderBlock2D"
]
}
local_repo/vae_decoder/config.json
-31
View File
@@ -1,31 +0,0 @@
{
"_class_name": "AutoencoderKL",
"_diffusers_version": "0.19.0.dev0",
"act_fn": "silu",
"block_out_channels": [
128,
256,
512,
512
],
"down_block_types": [
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D"
],
"force_upcast": true,
"in_channels": 3,
"latent_channels": 4,
"layers_per_block": 2,
"norm_num_groups": 32,
"out_channels": 3,
"sample_size": 1024,
"scaling_factor": 0.13025,
"up_block_types": [
"UpDecoderBlock2D",
"UpDecoderBlock2D",
"UpDecoderBlock2D",
"UpDecoderBlock2D"
]
}
local_repo/vae_encoder/config.json
-31
View File
@@ -1,31 +0,0 @@
{
"_class_name": "AutoencoderKL",
"_diffusers_version": "0.19.0.dev0",
"act_fn": "silu",
"block_out_channels": [
128,
256,
512,
512
],
"down_block_types": [
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D"
],
"force_upcast": true,
"in_channels": 3,
"latent_channels": 4,
"layers_per_block": 2,
"norm_num_groups": 32,
"out_channels": 3,
"sample_size": 1024,
"scaling_factor": 0.13025,
"up_block_types": [
"UpDecoderBlock2D",
"UpDecoderBlock2D",
"UpDecoderBlock2D",
"UpDecoderBlock2D"
]
}
msdiffusion/__init__.py
-1
View File
@@ -1 +0,0 @@
msdiffusion/dataset/__init__.py
-1
View File
@@ -1 +0,0 @@
msdiffusion/dataset/prompts/msbench.py
-132
View File
@@ -1,132 +0,0 @@
data_class_dict = {
"backpack": "backpack",
"cat": "cat",
"cat2": "cat",
"clock": "clock",
"colorful_sneaker": "sneaker",
"dog": "dog",
"dog2": "dog",
"dog5": "dog",
"dog6": "dog",
"dog7": "dog",
"dog8": "dog",
"duck_toy": "toy",
"grey_sloth_plushie": "stuffed animal",
"teapot": "teapot",
"wolf_plushie": "stuffed animal",
"cat1": "cat",
"dog1": "dog",
"lantern": "lantern",
"flower": "flower",
"teddybear": "stuffed animal",
"forest": "forest",
"mountain": "mountain",
"barn": "barn",
"waterfall": "waterfall",
"cap": "cap",
"cap1": "cap",
"hat": "hat",
"hat1": "hat",
"hat2": "hat",
"dress": "dress",
"dress1": "dress",
"shorts1": "shorts",
"trousers": "trousers",
"jeans": "jeans",
"coat": "coat",
"jacket": "jacket",
"t-shirt": "t-shirt",
"shirt": "shirt",
"shirt1": "shirt",
"sweater": "sweater",
}
two_subject_combinations = ["living_living", "living_object", "object_object", "living_upwearing", "living_midwearing",
"living_wholewearing", "midwearing_downwearing", "living_scene", "object_scene"]
three_subject_combinations = ["living_living_living", "object_object_object", "living_object_scene",
"upwearing_midwearing_downwearing"]
subject_combination_boxes = {
"living_living": [[0., 0.25, 0.5, 0.75], [0.5, 0.25, 1., 0.75]],
"living_object": [[0., 0.25, 0.5, 0.75], [0.5, 0.25, 1., 0.75]],
"object_object": [[0., 0.25, 0.5, 0.75], [0.5, 0.25, 1., 0.75]],
"living_upwearing": [[0.25, 0.25, 0.75, 0.75], [0.25, 0., 0.75, 0.25]],
"living_midwearing": [[0.25, 0.25, 0.75, 0.75], [0.25, 0.25, 0.75, 0.75]],
"living_wholewearing": [[0.25, 0.25, 0.75, 0.75], [0.25, 0.25, 0.75, 1.]],
"midwearing_downwearing": [[0.25, 0.25, 0.75, 0.6], [0.25, 0.6, 0.75, 1.]],
"living_scene": [[0.25, 0.25, 0.75, 0.75], [0., 0., 1., 1.]],
"object_scene": [[0.25, 0.25, 0.75, 0.75], [0., 0., 1., 1.]],
"living_living_living": [[0., 0.25, 0.35, 0.75], [0.35, 0.25, 0.65, 0.75], [0.65, 0.25, 1., 0.75]],
"object_object_object": [[0., 0.25, 0.35, 0.75], [0.35, 0.25, 0.65, 0.75], [0.65, 0.25, 1., 0.75]],
"living_object_scene": [[0., 0.25, 0.5, 0.75], [0.5, 0.25, 1., 0.75], [0., 0., 1., 1.]],
"upwearing_midwearing_downwearing": [[0.25, 0., 0.75, 0.25], [0.25, 0.25, 0.75, 0.6], [0.25, 0.6, 0.75, 1.]],
}
normal_prompts_two_subjects = [
"a {0} and a {1} in a room",
"a {0} and a {1} in the snow",
"a {0} and a {1} in the jungle",
"a {0} and a {1} on the beach",
"a {0} and a {1} on the grass",
"a {0} and a {1} on a cobblestone street",
]
normal_prompts_three_subjects = [
"a {0}, a {1}, and a {2} in a room",
"a {0}, a {1}, and a {2} in the snow",
"a {0}, a {1}, and a {2} in the jungle",
"a {0}, a {1}, and a {2} on the beach",
"a {0}, a {1}, and a {2} on the grass",
"a {0}, a {1}, and a {2} on a cobblestone street",
]
scene_prompts_two_subjects = [
"a {0} with a {1} in the background",
]
scene_prompts_three_subjects = [
"a {0} and {1} with a {2} in the background",
]
wearing_prompts_two_subjects_0 = [
"a {0} wearing a {1} in a room",
"a {0} wearing a {1} in the snow",
"a {0} wearing a {1} in the jungle",
"a {0} wearing a {1} on the beach",
"a {0} wearing a {1} on the grass",
"a {0} wearing a {1} on a cobblestone street",
]
wearing_prompts_two_subjects_1 = [
"a woman wearing a {0} and {1} in a room",
"a woman wearing a {0} and {1} in the snow",
"a woman wearing a {0} and {1} in the jungle",
"a woman wearing a {0} and {1} on the beach",
"a woman wearing a {0} and {1} on the grass",
"a woman wearing a {0} and {1} on a cobblestone street"
]
wearing_prompts_three_subjects = [
"a woman wearing a {0}, a {1}, and a {2} in a room",
"a woman wearing a {0}, a {1}, and a {2} in the snow",
"a woman wearing a {0}, a {1}, and a {2} in the jungle",
"a woman wearing a {0}, a {1}, and a {2} on the beach",
"a woman wearing a {0}, a {1}, and a {2} on the grass",
"a woman wearing a {0}, a {1}, and a {2} on a cobblestone street"
]
subject_combination_prompts = {
"living_living": normal_prompts_two_subjects,
"living_object": normal_prompts_two_subjects,
"object_object": normal_prompts_two_subjects,
"living_upwearing": wearing_prompts_two_subjects_0,
"living_midwearing": wearing_prompts_two_subjects_0,
"living_wholewearing": wearing_prompts_two_subjects_0,
"midwearing_downwearing": wearing_prompts_two_subjects_1,
"living_scene": scene_prompts_two_subjects,
"object_scene": scene_prompts_two_subjects,
"living_living_living": normal_prompts_three_subjects,
"object_object_object": normal_prompts_three_subjects,
"living_object_scene": scene_prompts_three_subjects,
"upwearing_midwearing_downwearing": wearing_prompts_three_subjects,
}
msdiffusion/models/__init__.py
-5
View File
@@ -1,5 +0,0 @@
from .model import MSAdapter
__all__ = [
"MSAdapter",
]
msdiffusion/models/attention_processor.py
-494
View File
@@ -1,494 +0,0 @@
# modified from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py
# and https://github.com/tencent-ailab/IP-Adapter/blob/main/ip_adapter/attention_processor.py
import torch
import torch.nn as nn
import torch.nn.functional as F
def minmax_normalize(batch_maps):
min_val = batch_maps.min(dim=-1, keepdim=True)[0].min(dim=-2, keepdim=True)[0]
max_val = batch_maps.max(dim=-1, keepdim=True)[0].max(dim=-2, keepdim=True)[0]
return (batch_maps - min_val) / (max_val - min_val + 1e-5)
class AttnProcessor2_0(torch.nn.Module):
r"""
Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
"""
def __init__(
self,
hidden_size=None,
cross_attention_dim=None,
):
super().__init__()
if not hasattr(F, "scaled_dot_product_attention"):
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
def __call__(
self,
attn,
hidden_states,
encoder_hidden_states=None,
attention_mask=None,
temb=None,
boxes=None,
phrase_idxes=None,
eot_idxes=None,
):
residual = hidden_states
if attn.spatial_norm is not None:
hidden_states = attn.spatial_norm(hidden_states, temb)
input_ndim = hidden_states.ndim
if input_ndim == 4:
batch_size, channel, height, width = hidden_states.shape
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
batch_size, sequence_length, _ = (
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
)
if attention_mask is not None:
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
# scaled_dot_product_attention expects attention_mask shape to be
# (batch, heads, source_length, target_length)
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
if attn.group_norm is not None:
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
query = attn.to_q(hidden_states)
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
elif attn.norm_cross:
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
key = attn.to_k(encoder_hidden_states)
value = attn.to_v(encoder_hidden_states)
inner_dim = key.shape[-1]
head_dim = inner_dim // attn.heads
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
# the output of sdp = (batch, num_heads, seq_len, head_dim)
# TODO: add support for attn.scale when we move to Torch 2.1
hidden_states = F.scaled_dot_product_attention(
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
)
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
hidden_states = hidden_states.to(query.dtype)
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
if input_ndim == 4:
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
if attn.residual_connection:
hidden_states = hidden_states + residual
hidden_states = hidden_states / attn.rescale_output_factor
return hidden_states
class MaskedIPAttnProcessor2_0(nn.Module):
def __init__(self, hidden_size, cross_attention_dim=None, scale=1.0, num_tokens=4, text_tokens=77,
need_text_attention_map=False, need_image_attention_map=True, num_dummy_tokens=4, mask_threshold=0.5,
use_psuedo_attention_mask=False, subject_scales=None, start_step=5):
super().__init__()
if not hasattr(F, "scaled_dot_product_attention"):
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
self.hidden_size = hidden_size
self.cross_attention_dim = cross_attention_dim
self.scale = scale
self.num_tokens = num_tokens
self.text_tokens = text_tokens
self.num_dummy_tokens = num_dummy_tokens
self.mask_threshold = mask_threshold
self.subject_scales = subject_scales
self.start_step = start_step
self.to_k_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
self.to_v_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
self.need_text_attention_map = need_text_attention_map
self.need_image_attention_map = need_image_attention_map
self.use_psuedo_attention_mask = use_psuedo_attention_mask
self.attention_maps = []
def prepare_attention_mask_qk(self, boxes, phrase_idxes, sequence_length_q, sequence_length_k,
batch_size, head_size, dtype, device, use_masked_text_attention=False):
if boxes is None:
return None, None
# TODO: only support square image now
num_patches_per_row = int(sequence_length_q ** 0.5)
box_idxes_start = torch.floor(boxes[:, :, 0:2] * num_patches_per_row)
box_idxes_end = torch.ceil(boxes[:, :, 2:4] * num_patches_per_row)
box_idxes = torch.cat([box_idxes_start, box_idxes_end], dim=-1)
box_masks = []
dummy_attention_mask = torch.ones((batch_size, sequence_length_q), dtype=dtype, device=device)
for box_idx in box_idxes.unbind(dim=1):
x_start_patch_idx, y_start_patch_idx, x_end_patch_idx, y_end_patch_idx = box_idx.unbind(dim=1)
x_indices = torch.arange(num_patches_per_row).unsqueeze(0).expand(batch_size, -1).to(device)
y_indices = torch.arange(num_patches_per_row).unsqueeze(0).expand(batch_size, -1).to(device)
x_mask = ((x_indices >= x_start_patch_idx.unsqueeze(1)) & (x_indices < x_end_patch_idx.unsqueeze(1))).to(dtype)
y_mask = ((y_indices >= y_start_patch_idx.unsqueeze(1)) & (y_indices < y_end_patch_idx.unsqueeze(1))).to(dtype)
box_mask = torch.bmm(y_mask.unsqueeze(2), x_mask.unsqueeze(1)).reshape(batch_size, -1)
box_masks.append(box_mask)
dummy_attention_mask = torch.clamp(dummy_attention_mask - box_mask, min=0)
# post mask
post_dummy_attention_mask = dummy_attention_mask.to(torch.bool)
post_dummy_attention_mask = post_dummy_attention_mask.repeat_interleave(head_size, dim=0)
attention_mask_qk_image = torch.stack(box_masks, dim=-1)
attention_mask_qk_image = attention_mask_qk_image.repeat_interleave(self.num_tokens, dim=-1)
attention_mask_qk_image = (1 - attention_mask_qk_image.to(dtype)) * -10000.0 # mask to bias
# use dummy image tokens to process the background
dummy_attention_mask = dummy_attention_mask.unsqueeze(-1).repeat_interleave(self.num_dummy_tokens, dim=-1)
dummy_attention_mask = (1 - dummy_attention_mask) * -10000.0
attention_mask_qk_image = torch.cat([dummy_attention_mask, attention_mask_qk_image], dim=-1)
if attention_mask_qk_image.shape[0] < batch_size*head_size:
attention_mask_qk_image = attention_mask_qk_image.repeat_interleave(head_size, dim=0)
if use_masked_text_attention:
attention_mask_qk_text = torch.ones((batch_size, sequence_length_q, sequence_length_k), dtype=dtype, device=device)
for i in range(batch_size):
for j in range(len(box_masks)):
start_idx, end_idx = int(phrase_idxes[i, j, 0].item()), int(phrase_idxes[i, j, 1].item())
if start_idx == 0 and end_idx == 0:
continue
attention_mask_qk_text[i, :, start_idx:end_idx] = box_masks[j][i, ...].unsqueeze(-1)
attention_mask_qk_text = (1 - attention_mask_qk_text) * -10000.0
if attention_mask_qk_text.shape[0] < batch_size*head_size:
attention_mask_qk_text = attention_mask_qk_text.repeat_interleave(head_size, dim=0)
else:
attention_mask_qk_text = None
return attention_mask_qk_image, attention_mask_qk_text, post_dummy_attention_mask
def get_text_attention_maps(self, attention_probs, boxes, phrase_idxes, head_size):
bsz = boxes.shape[0]
_, num_tokens_q, num_tokens_k = attention_probs.shape
attention_probs = attention_probs.view(bsz, head_size, num_tokens_q, num_tokens_k)
num_ref = boxes.shape[1]
h = w = int(num_tokens_q ** 0.5)
batch_attention_maps = []
for i in range(bsz):
sample_attention_maps = []
for j in range(num_ref):
start_idx, end_idx = int(phrase_idxes[i, j, 0].item()), int(phrase_idxes[i, j, 1].item())
if start_idx == 0 and end_idx == 0:
sample_attention_maps.append(
torch.zeros(num_tokens_q, dtype=attention_probs.dtype, device=attention_probs.device))
else:
attention_map = attention_probs[i, :, :,
start_idx:end_idx] # [num_heads, num_tokens_q, num_tokens_phrase]
attention_map = torch.mean(torch.mean(attention_map, dim=-1), dim=0) # [num_tokens_q]
sample_attention_maps.append(attention_map)
batch_attention_maps.append(torch.stack(sample_attention_maps))
self.attention_maps.append(torch.stack(batch_attention_maps).reshape(bsz, num_ref, h, w))
def get_psuedo_attention_mask(self, head_size):
# text_attention_maps = self.attention_maps[-1] # [bsz, num_ref, h, w]
if not self.use_psuedo_attention_mask or len(self.attention_maps) < self.start_step:
return None, None
text_attention_maps = torch.stack(self.attention_maps).mean(dim=0) # [bsz, num_ref, h, w]
text_attention_maps = minmax_normalize(text_attention_maps)
dtype, device = text_attention_maps.dtype, text_attention_maps.device
bsz, num_ref, h, w = text_attention_maps.shape
seq_len_q = h * w
text_attention_maps = text_attention_maps.view(bsz, num_ref, -1)
text_attention_maps = text_attention_maps.transpose(1, 2) # [bsz, h*w, num_ref]
# use threshold to get the mask
psuedo_attention_mask = (text_attention_maps > self.mask_threshold).to(dtype)
psuedo_dummy_attention_mask = torch.ones((bsz, seq_len_q), dtype=dtype, device=device)
for i in range(num_ref):
psuedo_box_mask = psuedo_attention_mask[..., i]
psuedo_dummy_attention_mask = torch.clamp(psuedo_dummy_attention_mask - psuedo_box_mask, min=0)
# post mask
post_psuedo_dummy_attention_mask = psuedo_dummy_attention_mask.to(torch.bool)
post_psuedo_dummy_attention_mask = post_psuedo_dummy_attention_mask.repeat_interleave(head_size, dim=0)
psuedo_attention_mask = psuedo_attention_mask.repeat_interleave(self.num_tokens, dim=-1)
psuedo_attention_mask = (1 - psuedo_attention_mask) * -10000.0 # mask to bias
psuedo_dummy_attention_mask = psuedo_dummy_attention_mask.unsqueeze(-1).repeat_interleave(self.num_dummy_tokens, dim=-1)
psuedo_dummy_attention_mask = (1 - psuedo_dummy_attention_mask) * -10000.0
psuedo_attention_mask = torch.cat([psuedo_dummy_attention_mask, psuedo_attention_mask], dim=-1)
if psuedo_attention_mask.shape[0] < bsz * head_size:
psuedo_attention_mask = psuedo_attention_mask.repeat_interleave(head_size, dim=0)
return psuedo_attention_mask, post_psuedo_dummy_attention_mask
def __call__(
self,
attn,
hidden_states,
encoder_hidden_states=None,
attention_mask=None,
temb=None,
boxes=None,
phrase_idxes=None,
eot_idxes=None,
):
residual = hidden_states
if attn.spatial_norm is not None:
hidden_states = attn.spatial_norm(hidden_states, temb)
input_ndim = hidden_states.ndim
if input_ndim == 4:
batch_size, channel, height, width = hidden_states.shape
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
batch_size, sequence_length, _ = (
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
)
if attention_mask is not None:
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
# scaled_dot_product_attention expects attention_mask shape to be
# (batch, heads, source_length, target_length)
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
rf_attention_mask = None
custom_attention_masks = self.prepare_attention_mask_qk(boxes, phrase_idxes, hidden_states.shape[1],
self.text_tokens, batch_size, attn.heads,
hidden_states.dtype, hidden_states.device,
use_masked_text_attention=False)
attention_mask_qk_image, attention_mask_qk_text, dummy_attention_mask = custom_attention_masks
if attention_mask_qk_image is not None:
attention_mask_qk_image = attention_mask_qk_image.view(batch_size, attn.heads, -1, attention_mask_qk_image.shape[-1])
if attention_mask_qk_text is not None:
attention_mask_qk_text = attention_mask_qk_text.view(batch_size, attn.heads, -1, attention_mask_qk_text.shape[-1])
if attn.group_norm is not None:
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
query = attn.to_q(hidden_states)
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
else:
# get encoder_hidden_states, ip_hidden_states
# end_pos = encoder_hidden_states.shape[1] - self.num_tokens
end_pos = self.text_tokens
encoder_hidden_states, ip_hidden_states = (
encoder_hidden_states[:, :end_pos, :],
encoder_hidden_states[:, end_pos:, :],
)
attention_mask, rf_attention_mask = (
attention_mask[:, :, :, :end_pos],
attention_mask[:, :, :, end_pos:],
) if attention_mask is not None else (None, None)
if attn.norm_cross:
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
key = attn.to_k(encoder_hidden_states)
value = attn.to_v(encoder_hidden_states)
inner_dim = key.shape[-1]
head_dim = inner_dim // attn.heads
attention_mask = attention_mask_qk_text if attention_mask_qk_text is not None else attention_mask
if not self.need_text_attention_map:
# original attention 2.0
new_query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
# the output of sdp = (batch, num_heads, seq_len, head_dim)
hidden_states = F.scaled_dot_product_attention(
new_query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
)
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
hidden_states = hidden_states.to(query.dtype)
else:
# we need get the attention map, so use the previous attention
new_query = attn.head_to_batch_dim(query)
key = attn.head_to_batch_dim(key)
value = attn.head_to_batch_dim(value)
if attention_mask is not None:
attention_mask = attention_mask.view(batch_size*attn.heads, -1, attention_mask.shape[-1])
attention_probs = attn.get_attention_scores(new_query, key, attention_mask)
self.get_text_attention_maps(attention_probs, boxes, phrase_idxes, attn.heads)
hidden_states = torch.bmm(attention_probs, value)
hidden_states = attn.batch_to_head_dim(hidden_states)
# get psuedo attention mask for image: better start after some timesteps
psuedo_attention_mask, psuedo_dummy_attention_mask = self.get_psuedo_attention_mask(attn.heads)
if psuedo_attention_mask is not None:
psuedo_attention_mask = psuedo_attention_mask.view(batch_size, attn.heads, -1,
psuedo_attention_mask.shape[-1])
ip_key = self.to_k_ip(ip_hidden_states)
ip_value = self.to_v_ip(ip_hidden_states)
rf_attention_mask = attention_mask_qk_image if attention_mask_qk_image is not None else rf_attention_mask
rf_attention_mask = psuedo_attention_mask if psuedo_attention_mask is not None else rf_attention_mask
dummy_attention_mask = psuedo_dummy_attention_mask if psuedo_dummy_attention_mask is not None else dummy_attention_mask
if not self.need_image_attention_map:
new_query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
# the output of sdp = (batch, num_heads, seq_len, head_dim)
ip_hidden_states = F.scaled_dot_product_attention(
new_query, ip_key, ip_value, attn_mask=rf_attention_mask, dropout_p=0.0, is_causal=False
)
ip_hidden_states = ip_hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
ip_hidden_states = ip_hidden_states.to(query.dtype)
else:
new_query = attn.head_to_batch_dim(query)
ip_key = attn.head_to_batch_dim(ip_key)
ip_value = attn.head_to_batch_dim(ip_value)
if rf_attention_mask is not None:
rf_attention_mask = rf_attention_mask.view(batch_size*attn.heads, -1, rf_attention_mask.shape[-1])
ip_attention_probs = attn.get_attention_scores(new_query, ip_key, rf_attention_mask)
# mask attention_probs in background
ip_attention_probs = torch.where(dummy_attention_mask.unsqueeze(-1), torch.zeros_like(ip_attention_probs), ip_attention_probs)
if self.subject_scales is not None:
# apply different scales to different subjects
subject_scales = torch.tensor(self.subject_scales, dtype=ip_attention_probs.dtype, device=ip_attention_probs.device)
subject_scales = subject_scales.unsqueeze(0).unsqueeze(0).repeat_interleave(self.num_tokens, dim=-1)
dummy_subject_scales = torch.ones((1, 1, 1), dtype=ip_attention_probs.dtype, device=ip_attention_probs.device).repeat_interleave(self.num_dummy_tokens, dim=-1)
subject_scales = torch.cat([dummy_subject_scales, subject_scales], dim=-1)
ip_attention_probs = ip_attention_probs * subject_scales
ip_hidden_states = torch.bmm(ip_attention_probs, ip_value)
ip_hidden_states = attn.batch_to_head_dim(ip_hidden_states)
if self.subject_scales is None:
hidden_states = hidden_states + self.scale * ip_hidden_states
else:
hidden_states = hidden_states + ip_hidden_states
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
if input_ndim == 4:
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
if attn.residual_connection:
hidden_states = hidden_states + residual
hidden_states = hidden_states / attn.rescale_output_factor
return hidden_states
class CNAttnProcessor2_0:
r"""
Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
"""
def __init__(self, num_tokens=4, text_tokens=77):
if not hasattr(F, "scaled_dot_product_attention"):
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
self.num_tokens = num_tokens
self.text_tokens = text_tokens
def __call__(
self,
attn,
hidden_states,
encoder_hidden_states=None,
attention_mask=None,
temb=None,
):
residual = hidden_states
if attn.spatial_norm is not None:
hidden_states = attn.spatial_norm(hidden_states, temb)
input_ndim = hidden_states.ndim
if input_ndim == 4:
batch_size, channel, height, width = hidden_states.shape
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
batch_size, sequence_length, _ = (
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
)
if attention_mask is not None:
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
# scaled_dot_product_attention expects attention_mask shape to be
# (batch, heads, source_length, target_length)
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
rf_attention_mask = None
if attn.group_norm is not None:
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
query = attn.to_q(hidden_states)
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
else:
# end_pos = encoder_hidden_states.shape[1] - self.num_tokens
end_pos = self.text_tokens
encoder_hidden_states = encoder_hidden_states[:, :end_pos] # only use text
attention_mask = attention_mask[:, :, :end_pos]
if attn.norm_cross:
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
key = attn.to_k(encoder_hidden_states)
value = attn.to_v(encoder_hidden_states)
inner_dim = key.shape[-1]
head_dim = inner_dim // attn.heads
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
# the output of sdp = (batch, num_heads, seq_len, head_dim)
# TODO: add support for attn.scale when we move to Torch 2.1
hidden_states = F.scaled_dot_product_attention(
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
)
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
hidden_states = hidden_states.to(query.dtype)
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
if input_ndim == 4:
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
if attn.residual_connection:
hidden_states = hidden_states + residual
hidden_states = hidden_states / attn.rescale_output_factor
return hidden_states
msdiffusion/models/model.py
-263
View File
@@ -1,263 +0,0 @@
import os
import math
from typing import List
import torch
import torch.nn as nn
from diffusers.pipelines.controlnet import MultiControlNetModel
from .attention_processor import MaskedIPAttnProcessor2_0 as IPAttnProcessor, AttnProcessor2_0 as AttnProcessor, \
CNAttnProcessor2_0 as CNAttnProcessor
class MSAdapter(torch.nn.Module):
def __init__(self, unet,image_proj_model, adapter_modules=None, ckpt_path=None,
num_tokens=4, text_tokens=77, max_rn=4, num_dummy_tokens=4,
device="cuda", controlnet=None):
super().__init__()
self.unet = unet
self.image_proj_model = image_proj_model
self.adapter_modules = adapter_modules
self.num_tokens = num_tokens
self.num_dummy_tokens = num_dummy_tokens
self.text_tokens = text_tokens
self.max_rn = max_rn
self.device = device
self.controlnet = controlnet
self.cross_attention_dim = self.unet.config.cross_attention_dim
# set attention processor when inference
if self.adapter_modules is None:
self.set_ms_adapter()
# dummy image tokens
self.dummy_image_tokens = nn.Parameter(torch.randn(1, self.num_dummy_tokens, self.cross_attention_dim))
if ckpt_path is not None:
self.load_from_checkpoint(ckpt_path)
def load_from_checkpoint(self, ckpt_path: str):
if os.path.isdir(ckpt_path):
ckpt_path = os.path.join(ckpt_path, "ms_adapter.bin")
# Calculate original checksums
orig_ip_proj_sum = torch.sum(torch.stack([torch.sum(p) for p in self.image_proj_model.parameters()]))
orig_adapter_sum = torch.sum(torch.stack([torch.sum(p) for p in self.adapter_modules.parameters()]))
state_dict = torch.load(ckpt_path, map_location="cpu")
# Load state dict for image_proj_model and adapter_modules when using resampler
self.image_proj_model.load_state_dict(state_dict["image_proj"], strict=False)
self.adapter_modules.load_state_dict(state_dict["ms_adapter"], strict=False)
self.load_state_dict({"dummy_image_tokens": state_dict["dummy_image_tokens"]}, strict=False)
# Calculate new checksums
new_ip_proj_sum = torch.sum(torch.stack([torch.sum(p) for p in self.image_proj_model.parameters()]))
new_adapter_sum = torch.sum(torch.stack([torch.sum(p) for p in self.adapter_modules.parameters()]))
# Verify if the weights have changed
assert orig_ip_proj_sum != new_ip_proj_sum, "Weights of image_proj_model did not change!"
assert orig_adapter_sum != new_adapter_sum, "Weights of adapter_modules did not change!"
print(f"Successfully loaded weights from checkpoint {ckpt_path}")
def set_ms_adapter(self, weight_dtype=torch.float16, cache_attention_maps=True):
# set attention processor
attn_procs = {}
for name in self.unet.attn_processors.keys():
cross_attention_dim = None if name.endswith("attn1.processor") else self.unet.config.cross_attention_dim
if name.startswith("mid_block"):
hidden_size = self.unet.config.block_out_channels[-1]
elif name.startswith("up_blocks"):
block_id = int(name[len("up_blocks.")])
hidden_size = list(reversed(self.unet.config.block_out_channels))[block_id]
elif name.startswith("down_blocks"):
block_id = int(name[len("down_blocks.")])
hidden_size = self.unet.config.block_out_channels[block_id]
if cross_attention_dim is None:
attn_procs[name] = AttnProcessor()
else:
attn_procs[name] = IPAttnProcessor(
hidden_size=hidden_size,
cross_attention_dim=cross_attention_dim,
num_tokens=self.num_tokens,
text_tokens=self.text_tokens,
).to(self.device, dtype=weight_dtype)
self.unet.set_attn_processor(attn_procs)
self.adapter_modules = torch.nn.ModuleList(self.unet.attn_processors.values())
if self.controlnet is not None:
if isinstance(self.controlnet, MultiControlNetModel):
for controlnet in self.controlnet.nets:
controlnet.set_attn_processor(
CNAttnProcessor(text_tokens=self.text_tokens, num_tokens=self.num_tokens))
else:
self.controlnet.set_attn_processor(
CNAttnProcessor(text_tokens=self.text_tokens, num_tokens=self.num_tokens))
@torch.inference_mode()
def get_image_embeds(self, processed_images, image_encoder=None, image_proj_type="linear",
image_encoder_type="clip", weight_dtype=torch.float16,):
# get image embeds
#processed_images: [bsz, rn, ...]
# if use_repo:
# processed_images = processed_images.view(-1, processed_images.shape[-3], processed_images.shape[-2],
# processed_images.shape[-1]) # (bsz*rn, ...)
# #print(processed_images.shape) # torch.Size([1, 3, 224, 224])
if image_proj_type == "resampler":
# if use_repo:
# image_embeds = image_encoder(processed_images.to(self.device, dtype=weight_dtype),
# output_hidden_states=True).hidden_states[
# -2] # (bsz*rn, num_tokens, embedding_dim)
# else:
image_encoder = image_encoder.encode_image
processed_images.to(self.device, dtype=weight_dtype)
#image_embeds = image_encoder(processed_images)["image_embeds"]
#image_embeds = image_encoder(processed_images)["last_hidden_state"]
image_embeds=image_encoder(processed_images)["penultimate_hidden_states"]
# print(image_embeds.shape)
else:
image_embeds = image_encoder(
processed_images.to(self.device, dtype=weight_dtype)).image_embeds # (bsz*rn, embedding_dim)
return image_embeds # [bsz*rn, ...]
def set_scale(self, scale, subject_scales):
for attn_processor in self.pipe.unet.attn_processors.values():
if isinstance(attn_processor, IPAttnProcessor):
attn_processor.scale = scale
attn_processor.subject_scales = subject_scales
def enable_psuedo_attention_mask(self, mask_threshold=0.5, start_step=5):
for attn_processor in self.pipe.unet.attn_processors.values():
if isinstance(attn_processor, IPAttnProcessor):
attn_processor.mask_threshold = mask_threshold
attn_processor.start_step = start_step
attn_processor.use_psuedo_attention_mask = True
attn_processor.need_text_attention_map = True
attn_processor.attention_maps = [] # clear attention maps
def generate(self, pipe, pil_images=None, processed_images=None, prompt=None, negative_prompt=None, scale=1.0,
num_samples=4, seed=None, guidance_scale=7.5, num_inference_steps=50, image_processor=None,
image_encoder=None, image_proj_type="linear", image_encoder_type="clip", weight_dtype=torch.float16,
boxes=None, phrases=None, drop_grounding_tokens=None, phrase_idxes=None,
eot_idxes=None, height=1024, width=1024, subject_scales=None, mask_threshold=None, start_step=5,
**kwargs):
# generate images (validation&inference)
self.pipe = pipe
self.set_scale(scale, subject_scales)
if mask_threshold is not None:
self.enable_psuedo_attention_mask(mask_threshold, start_step)
# pil_images: [[xxx, xxx, xxx], [xxx, xxx, xxx], ...]
bsz = len(pil_images) # only support bsz=1 now
if processed_images is None:
# write in this way to promise it can be extended to batch in the future
processed_images = []
for pil_image in pil_images:
processed_image = image_processor(images=pil_image, return_tensors="pt").pixel_values
processed_images.append(processed_image)
processed_images = torch.stack(processed_images, dim=0)
num_prompts = bsz
if prompt is None:
prompt = "best quality, high quality"
if negative_prompt is None:
negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality" # duplicate
if not isinstance(prompt, List):
prompt = [prompt] * num_prompts
if not isinstance(negative_prompt, List):
negative_prompt = [negative_prompt] * num_prompts
cross_attention_kwargs = None
grounding_kwargs = None
if boxes is not None:
boxes = torch.tensor(boxes).to(self.device, weight_dtype)
if phrases is not None:
drop_grounding_tokens = drop_grounding_tokens if drop_grounding_tokens is not None else [0] * bsz
batch_boxes = boxes.view(bsz * boxes.shape[1], -1)
# write in this way to promise it can be extended to batch in the future
phrase_input_ids = []
for phrase in phrases:
phrase_input_id = pipe.tokenizer(phrase, max_length=pipe.tokenizer.model_max_length,
padding="max_length", truncation=True,
return_tensors="pt").input_ids
phrase_input_ids.append(phrase_input_id)
phrase_input_ids = torch.stack(phrase_input_ids)
phrase_input_ids = phrase_input_ids.view(-1, phrase_input_ids.shape[-1])
phrase_embeds = pipe.text_encoder(phrase_input_ids.to(self.device)).pooler_output
#print(phrase_embeds.shape,batch_boxes.shape)#torch.Size([1, 768]) torch.Size([1, 4])
grounding_kwargs = {"boxes": batch_boxes, "phrase_embeds": phrase_embeds,
"drop_grounding_tokens": drop_grounding_tokens}
else:
grounding_kwargs = None
boxes = torch.repeat_interleave(boxes, repeats=num_samples, dim=0)
uncond_boxes = torch.zeros_like(boxes)
boxes = torch.cat([uncond_boxes, boxes], dim=0)
cross_attention_kwargs = {"boxes": boxes}
if phrase_idxes is not None:
phrase_idxes = torch.tensor(phrase_idxes).to(self.device, torch.int)
eot_idxes = torch.tensor(eot_idxes).to(self.device, torch.int)
phrase_idxes = torch.repeat_interleave(phrase_idxes, repeats=num_samples, dim=0)
eot_idxes = torch.repeat_interleave(eot_idxes, repeats=num_samples, dim=0)
uncond_phrase_idxes = torch.zeros_like(phrase_idxes)
uncond_eot_idxes = torch.zeros_like(eot_idxes)
phrase_idxes = torch.cat([uncond_phrase_idxes, phrase_idxes], dim=0)
eot_idxes = torch.cat([uncond_eot_idxes, eot_idxes], dim=0)
if cross_attention_kwargs is None:
cross_attention_kwargs = {"phrase_idxes": phrase_idxes, "eot_idxes": eot_idxes}
else:
cross_attention_kwargs["phrase_idxes"] = phrase_idxes
cross_attention_kwargs["eot_idxes"] = eot_idxes
with (torch.inference_mode()):
image_embeds = self.get_image_embeds(processed_images, image_encoder, image_proj_type=image_proj_type,
image_encoder_type=image_encoder_type, weight_dtype=weight_dtype,)
#print(image_embeds.shape) #torch.Size([1, 257, 1664])
image_embeds=image_embeds.clone().detach().to(self.device, weight_dtype)
image_prompt_embeds = self.image_proj_model(image_embeds, grounding_kwargs=grounding_kwargs)
del image_embeds
image_prompt_embeds = image_prompt_embeds.view(bsz, -1, image_prompt_embeds.shape[-2],
image_prompt_embeds.shape[
-1]) # (bsz, rn, num_tokens, cross_attention_dim)
image_prompt_embeds = image_prompt_embeds.view(bsz,
image_prompt_embeds.shape[-3] * image_prompt_embeds.shape[
-2],
image_prompt_embeds.shape[
-1]) # (bsz, total_num_tokens*rn, cross_attention_dim)
image_prompt_embeds = torch.cat([self.dummy_image_tokens, image_prompt_embeds], dim=1)
uncond_image_prompt_embeds = torch.zeros_like(image_prompt_embeds)
bs_embed, seq_len, _ = image_prompt_embeds.shape
image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1)
image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1)
uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1)
prompt_embeds_, negative_prompt_embeds_, pooled_prompt_embeds, negative_pooled_prompt_embeds = pipe.encode_prompt(
prompt,
device=self.device,
num_images_per_prompt=num_samples,
do_classifier_free_guidance=True,
negative_prompt=negative_prompt,
)
prompt_embeds = torch.cat([prompt_embeds_, image_prompt_embeds], dim=1)
negative_prompt_embeds = torch.cat([negative_prompt_embeds_, uncond_image_prompt_embeds], dim=1)
generator = torch.Generator(self.device).manual_seed(seed) if seed is not None else None
images = pipe(
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
generator=generator,
cross_attention_kwargs=cross_attention_kwargs,
height=height,
width=width,
**kwargs,
).images
return images
msdiffusion/models/projection.py
-257
View File
@@ -1,257 +0,0 @@
# modified from https://github.com/mlfoundations/open_flamingo/blob/main/open_flamingo/src/helpers.py
# and https://github.com/lucidrains/imagen-pytorch/blob/main/imagen_pytorch/imagen_pytorch.py
# and https://github.com/tencent-ailab/IP-Adapter/blob/main/ip_adapter/ip_adapter.py
# and https://github.com/tencent-ailab/IP-Adapter/blob/main/ip_adapter/resampler.py
import math
import torch
import torch.nn as nn
from einops import rearrange
from einops.layers.torch import Rearrange
class FourierEmbedder(nn.Module):
def __init__(self, num_freqs=64, temperature=100):
super().__init__()
self.num_freqs = num_freqs
self.temperature = temperature
freq_bands = temperature ** (torch.arange(num_freqs) / num_freqs)
freq_bands = freq_bands[None, None]
self.register_buffer("freq_bands", freq_bands, persistent=False)
def __call__(self, x):
x = self.freq_bands * x.unsqueeze(-1)
return torch.stack((x.sin(), x.cos()), dim=-1).permute(0, 2, 3, 1).reshape(x.shape[0], -1)
class ImageProjModel(torch.nn.Module):
"""Projection Model"""
def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4):
super().__init__()
self.cross_attention_dim = cross_attention_dim
self.clip_extra_context_tokens = clip_extra_context_tokens
self.proj = torch.nn.Linear(clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim)
self.norm = torch.nn.LayerNorm(cross_attention_dim)
def forward(self, image_embeds):
embeds = image_embeds
clip_extra_context_tokens = self.proj(embeds).reshape(
-1, self.clip_extra_context_tokens, self.cross_attention_dim
)
clip_extra_context_tokens = self.norm(clip_extra_context_tokens)
return clip_extra_context_tokens
# FFN
def FeedForward(dim, mult=4):
inner_dim = int(dim * mult)
return nn.Sequential(
nn.LayerNorm(dim),
nn.Linear(dim, inner_dim, bias=False),
nn.GELU(),
nn.Linear(inner_dim, dim, bias=False),
# nn.LayerNorm(dim),
)
def reshape_tensor(x, heads):
bs, length, width = x.shape
# (bs, length, width) --> (bs, length, n_heads, dim_per_head)
x = x.view(bs, length, heads, -1)
# (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
x = x.transpose(1, 2)
# (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head)
x = x.reshape(bs, heads, length, -1)
return x
class PerceiverAttention(nn.Module):
def __init__(self, *, dim, dim_head=64, heads=8):
super().__init__()
self.scale = dim_head**-0.5
self.dim_head = dim_head
self.heads = heads
inner_dim = dim_head * heads
self.norm1 = nn.LayerNorm(dim)
self.norm2 = nn.LayerNorm(dim)
self.to_q = nn.Linear(dim, inner_dim, bias=False)
self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
self.to_out = nn.Linear(inner_dim, dim, bias=False)
def forward(self, x, latents):
"""
Args:
x (torch.Tensor): image features
shape (b, n1, D)
latent (torch.Tensor): latent features
shape (b, n2, D)
"""
x = self.norm1(x)
latents = self.norm2(latents)
b, l, _ = latents.shape
q = self.to_q(latents)
kv_input = torch.cat((x, latents), dim=-2)
k, v = self.to_kv(kv_input).chunk(2, dim=-1)
q = reshape_tensor(q, self.heads)
k = reshape_tensor(k, self.heads)
v = reshape_tensor(v, self.heads)
# attention
scale = 1 / math.sqrt(math.sqrt(self.dim_head))
weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards
weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
out = weight @ v
out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
return self.to_out(out)
class Resampler(nn.Module):
def __init__(
self,
dim=1024,
depth=8,
dim_head=64,
heads=16,
num_queries=8,
embedding_dim=768,
output_dim=1024,
ff_mult=4,
max_seq_len: int = 257, # CLIP tokens + CLS token
apply_pos_emb: bool = False,
num_latents_mean_pooled: int = 0, # number of latents derived from mean pooled representation of the sequence
latent_init_mode: str = "random",
phrase_embeddings_dim: int = 1024,
fourier_freqs: int = 8,
):
super().__init__()
self.num_queries = num_queries
self.grounding_token_num = self.num_queries
self.dim = dim
self.pos_emb = nn.Embedding(max_seq_len, embedding_dim) if apply_pos_emb else None
self.latent_init_mode = latent_init_mode
if latent_init_mode == "random":
self.latents = nn.Parameter(torch.randn(1, self.latents_token_num, dim) / dim**0.5)
self.fourier_embedder = None
self.latent_proj = None
self.latent_norm = None
elif latent_init_mode == "grounding":
self.latents = None
self.grounding_latents = nn.Parameter(torch.randn(1, self.grounding_token_num, dim) / dim ** 0.5)
self.fourier_embedder = FourierEmbedder(num_freqs=fourier_freqs)
grounding_embedding_dim = phrase_embeddings_dim + fourier_freqs * 2 * 4 # 2: sin/cos, 4: xyxy
self.latent_proj = torch.nn.Sequential(
torch.nn.Linear(grounding_embedding_dim, grounding_embedding_dim * 2),
torch.nn.GELU(),
torch.nn.Linear(grounding_embedding_dim * 2, dim * self.grounding_token_num),
)
self.latent_norm = nn.LayerNorm(dim)
else:
raise ValueError(f"Invalid latent_init_mode: {latent_init_mode}")
self.proj_in = nn.Linear(embedding_dim, dim)
self.attention_norm = nn.LayerNorm(dim)
self.proj_out = nn.Linear(dim, output_dim)
self.norm_out = nn.LayerNorm(output_dim)
self.to_latents_from_mean_pooled_seq = (
nn.Sequential(
nn.LayerNorm(dim),
nn.Linear(dim, dim * num_latents_mean_pooled),
Rearrange("b (n d) -> b n d", n=num_latents_mean_pooled),
)
if num_latents_mean_pooled > 0
else None
)
self.layers = nn.ModuleList([])
for _ in range(depth):
self.layers.append(
nn.ModuleList(
[
PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
FeedForward(dim=dim, mult=ff_mult),
]
)
)
def forward(self, x, grounding_kwargs=None, shortcut=False, scale=1.0):
if self.pos_emb is not None:
n, device = x.shape[1], x.device
pos_emb = self.pos_emb(torch.arange(n, device=device))
x = x + pos_emb
if self.latent_init_mode == "random":
latents = self.latents.repeat(x.size(0), 1, 1)
elif self.latent_init_mode == "grounding":
boxes = grounding_kwargs["boxes"]
phrase_embeds = grounding_kwargs["phrase_embeds"]
fourier_embeds = self.fourier_embedder(boxes)
grounding_embeds = torch.cat((phrase_embeds, fourier_embeds), dim=-1)
drop_grounding_tokens = grounding_kwargs["drop_grounding_tokens"]
num_ref = x.shape[0] // len(drop_grounding_tokens)
drop_grounding_tokens = [item for item in drop_grounding_tokens for _ in range(num_ref)]
latents = self.latent_proj(grounding_embeds)
latents = latents.view(-1, self.grounding_token_num, self.dim)
latents = self.latent_norm(latents)
# drop grounding tokens to learnable latents
drop_num = len([item for item in drop_grounding_tokens if item == 1])
if drop_num > 0:
latents_ = []
learnable_latents = self.grounding_latents.repeat(drop_num, 1, 1)
cur_idx = 0
for latent, drop_grounding_token in zip(latents, drop_grounding_tokens):
if drop_grounding_token == 1:
latent = learnable_latents[cur_idx]
cur_idx += 1
latents_.append(latent)
latents = torch.stack(latents_)
else:
raise ValueError(f"Invalid latent_init_mode: {self.latent_init_mode}")
x = self.proj_in(x)
if self.to_latents_from_mean_pooled_seq:
meanpooled_seq = masked_mean(x, dim=1, mask=torch.ones(x.shape[:2], device=x.device, dtype=torch.bool))
meanpooled_latents = self.to_latents_from_mean_pooled_seq(meanpooled_seq)
latents = torch.cat((meanpooled_latents, latents), dim=-2)
init_latents = latents
for attn, ff in self.layers:
latents = attn(x, latents) + latents
latents = ff(latents) + latents
latents = self.attention_norm(latents)
latents = self.proj_out(latents)
if shortcut:
latents = init_latents + latents * scale
return self.norm_out(latents)
def masked_mean(t, *, dim, mask=None):
if mask is None:
return t.mean(dim=dim)
denom = mask.sum(dim=dim, keepdim=True)
mask = rearrange(mask, "b n -> b n 1")
masked_t = t.masked_fill(~mask, 0.0)
return masked_t.sum(dim=dim) / denom.clamp(min=1e-5)
msdiffusion/utils.py
-50
View File
@@ -1,50 +0,0 @@
import re
def get_phrase_idx(tokenizer, phrase, prompt, get_last_word=False, num=0):
def is_equal_words(pr_words, ph_words):
if len(pr_words) != len(ph_words):
return False
for pr_word, ph_word in zip(pr_words, ph_words):
if "-"+ph_word not in pr_word and ph_word != re.sub(r'[.!?,:]$', '', pr_word):
return False
return True
phrase_words = phrase.split()
if len(phrase_words) == 0:
return [0, 0], None
if get_last_word:
phrase_words = phrase_words[-1:]
# prompt_words = re.findall(r'\b[\w\'-]+\b', prompt)
prompt_words = prompt.split()
start = 1
end = 0
res_words = phrase_words
for i in range(len(prompt_words)):
if is_equal_words(prompt_words[i:i+len(phrase_words)], phrase_words):
if num != 0:
# skip this one
num -= 1
continue
end = start
res_words = prompt_words[i:i+len(phrase_words)]
res_words = [re.sub(r'[.!?,:]$', '', w) for w in res_words]
prompt_words[i+len(phrase_words)-1] = res_words[-1] # remove the last punctuation
for j in range(i, i+len(phrase_words)):
end += len(tokenizer.encode(prompt_words[j])) - 2
break
else:
start += len(tokenizer.encode(prompt_words[i])) - 2
if end == 0:
return [0, 0], None
return [start, end], res_words
def get_eot_idx(tokenizer, prompt):
words = prompt.split()
start = 1
for w in words:
start += len(tokenizer.encode(w)) - 2
return start
msdiffusion_native/__init__.py
+4
View File
@@ -0,0 +1,4 @@
"""Native MS-Diffusion nodes: pure torch + ComfyUI core."""
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
msdiffusion_native/core.py
+546
View File
@@ -0,0 +1,546 @@
"""MS-Diffusion native core.
Pure torch + ComfyUI core APIs. No diffusers, no transformers, no einops,
no omegaconf, no PIL. Everything here is a line-by-line port of the upstream
MS-Diffusion inference path (D:/Downloads/MS-Diffusion-main):
msdiffusion/models/projection.py -> FourierEmbedder / Resampler
msdiffusion/models/attention_processor.py-> MaskedIPAttnProcessor2_0
msdiffusion/models/model.py -> MSAdapter.generate (conditioning assembly)
The SDXL UNet itself, text encoding, CLIP-Vision, KSampler and VAE are all
left to ComfyUI native nodes. This file only supplies the *patch* that turns
native attn2 into the MS-Diffusion dual-branch cross-attention.
Key semantic facts (verified against real weights on 2026-11-06):
*_ms_adapter.bin* has exactly 70 (to_k_ip, to_v_ip) pairs, indexed
1,3,...,139 in diffusers attn_processors order. That order is
down_blocks -> up_blocks -> mid_block (module *registration* order),
whereas ComfyUI *executes* input_blocks -> middle_block -> output_blocks.
The two must be permuted, see build_block_weights().
* diffusers and ComfyUI use the SAME head partition for SDXL:
inner_dim 640 -> 10 heads, 1280 -> 20 heads, head_dim 64.
"""
import math
import os
import uuid
import torch
import torch.nn as nn
import torch.nn.functional as F
NUM_TOKENS = 16 # image tokens per subject
NUM_DUMMY_TOKENS = 4 # learnable "background" image tokens
TEXT_TOKENS = 77 # upstream hard-codes this split point
VISION_DIM = 1664 # CLIP-ViT-bigG/14 penultimate hidden size
CROSS_DIM = 2048 # SDXL cross_attention_dim
PHRASE_DIM = 768 # CLIP-L pooled output
RESAMPLER_DIM = 1280
FOURIER_FREQS = 8
HEAD_DIM = 64 # SDXL attention head dim (both frameworks)
# ---------------------------------------------------------------------------
# projection.py port (einops removed; runtime never touches masked_mean/Rearrange)
# ---------------------------------------------------------------------------
class FourierEmbedder(nn.Module):
def __init__(self, num_freqs=64, temperature=100):
super().__init__()
self.num_freqs = num_freqs
self.temperature = temperature
freq_bands = temperature ** (torch.arange(num_freqs) / num_freqs)
self.register_buffer("freq_bands", freq_bands[None, None], persistent=False)
def __call__(self, x):
x = self.freq_bands * x.unsqueeze(-1)
return torch.stack((x.sin(), x.cos()), dim=-1).permute(0, 2, 3, 1).reshape(x.shape[0], -1)
def reshape_tensor(x, heads):
bs, length, width = x.shape
x = x.view(bs, length, heads, -1).transpose(1, 2)
return x.reshape(bs, heads, length, -1)
class PerceiverAttention(nn.Module):
def __init__(self, dim, dim_head=64, heads=8):
super().__init__()
self.scale = dim_head ** -0.5
self.dim_head = dim_head
self.heads = heads
inner_dim = dim_head * heads
self.norm1 = nn.LayerNorm(dim)
self.norm2 = nn.LayerNorm(dim)
self.to_q = nn.Linear(dim, inner_dim, bias=False)
self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
self.to_out = nn.Linear(inner_dim, dim, bias=False)
def forward(self, x, latents):
x = self.norm1(x)
latents = self.norm2(latents)
b, l, _ = latents.shape
q = self.to_q(latents)
kv_input = torch.cat((x, latents), dim=-2)
k, v = self.to_kv(kv_input).chunk(2, dim=-1)
q = reshape_tensor(q, self.heads)
k = reshape_tensor(k, self.heads)
v = reshape_tensor(v, self.heads)
# upstream stability trick: scale BOTH operands, float32 softmax
scale = 1 / math.sqrt(math.sqrt(self.dim_head))
weight = (q * scale) @ (k * scale).transpose(-2, -1)
weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
out = weight @ v
out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
return self.to_out(out)
def feed_forward(dim, mult=4):
inner_dim = int(dim * mult)
return nn.Sequential(
nn.LayerNorm(dim),
nn.Linear(dim, inner_dim, bias=False),
nn.GELU(),
nn.Linear(inner_dim, dim, bias=False),
)
class Resampler(nn.Module):
"""Grounding-mode Resampler, key-for-key identical to upstream projection.py."""
def __init__(self, dim=RESAMPLER_DIM, depth=4, dim_head=64, heads=20,
num_queries=NUM_TOKENS, embedding_dim=VISION_DIM,
output_dim=CROSS_DIM, ff_mult=4, max_seq_len=257,
phrase_embeddings_dim=PHRASE_DIM, fourier_freqs=FOURIER_FREQS):
super().__init__()
self.num_queries = num_queries
self.grounding_token_num = num_queries
self.dim = dim
self.latent_init_mode = "grounding"
self.grounding_latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim ** 0.5)
self.fourier_embedder = FourierEmbedder(num_freqs=fourier_freqs)
grounding_embedding_dim = phrase_embeddings_dim + fourier_freqs * 2 * 4
self.latent_proj = nn.Sequential(
nn.Linear(grounding_embedding_dim, grounding_embedding_dim * 2),
nn.GELU(),
nn.Linear(grounding_embedding_dim * 2, dim * num_queries),
)
self.latent_norm = nn.LayerNorm(dim)
self.proj_in = nn.Linear(embedding_dim, dim)
self.attention_norm = nn.LayerNorm(dim)
self.proj_out = nn.Linear(dim, output_dim)
self.norm_out = nn.LayerNorm(output_dim)
self.layers = nn.ModuleList([])
for _ in range(depth):
self.layers.append(nn.ModuleList([
PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
feed_forward(dim=dim, mult=ff_mult),
]))
def forward(self, x, grounding_kwargs=None, shortcut=False, scale=1.0):
boxes = grounding_kwargs["boxes"]
phrase_embeds = grounding_kwargs["phrase_embeds"]
fourier_embeds = self.fourier_embedder(boxes)
grounding_embeds = torch.cat((phrase_embeds, fourier_embeds), dim=-1)
drop_grounding_tokens = grounding_kwargs["drop_grounding_tokens"]
num_ref = x.shape[0] // len(drop_grounding_tokens)
drop_grounding_tokens = [item for item in drop_grounding_tokens for _ in range(num_ref)]
latents = self.latent_proj(grounding_embeds)
latents = latents.view(-1, self.grounding_token_num, self.dim)
latents = self.latent_norm(latents)
drop_num = len([item for item in drop_grounding_tokens if item == 1])
if drop_num > 0:
latents_ = []
learnable_latents = self.grounding_latents.repeat(drop_num, 1, 1)
cur_idx = 0
for latent, flag in zip(latents, drop_grounding_tokens):
if flag == 1:
latent = learnable_latents[cur_idx]
cur_idx += 1
latents_.append(latent)
latents = torch.stack(latents_)
x = self.proj_in(x)
init_latents = latents
for attn, ff in self.layers:
latents = attn(x, latents) + latents
latents = ff(latents) + latents
latents = self.attention_norm(latents)
latents = self.proj_out(latents)
if shortcut:
latents = init_latents + latents * scale
return self.norm_out(latents)
class IPProjection(nn.Module):
def __init__(self, hidden_size):
super().__init__()
self.to_k_ip = nn.Linear(CROSS_DIM, hidden_size, bias=False)
self.to_v_ip = nn.Linear(CROSS_DIM, hidden_size, bias=False)
class MSAdapter(nn.Module):
"""Self-contained adapter. Keys are 1:1 with ms_adapter.bin after index remap."""
def __init__(self, hidden_sizes):
super().__init__()
self.hidden_sizes = list(hidden_sizes)
self.image_proj = Resampler()
self.adapters = nn.ModuleList([IPProjection(h) for h in self.hidden_sizes])
self.dummy_image_tokens = nn.Parameter(torch.randn(1, NUM_DUMMY_TOKENS, CROSS_DIM))
self._placed = None
def load(self, path):
if os.path.isdir(path):
path = os.path.join(path, "ms_adapter.bin")
sd = torch.load(path, map_location="cpu", weights_only=True)
missing_top = {"image_proj", "ms_adapter", "dummy_image_tokens"} - set(sd)
if missing_top:
raise RuntimeError("MS-Diffusion checkpoint missing groups: %s" % sorted(missing_top))
self.image_proj.load_state_dict(sd["image_proj"], strict=True)
proj_sd = {}
for i in range(len(self.hidden_sizes)):
src = 2 * i + 1 # 1,3,...,139
for name in ("to_k_ip", "to_v_ip"):
proj_sd["%d.%s.weight" % (i, name)] = sd["ms_adapter"]["%d.%s.weight" % (src, name)]
self.adapters.load_state_dict(proj_sd, strict=True)
assert sd["dummy_image_tokens"].shape == self.dummy_image_tokens.shape
with torch.no_grad():
self.dummy_image_tokens.copy_(sd["dummy_image_tokens"])
for i, h in enumerate(self.hidden_sizes):
w = self.adapters[i].to_k_ip.weight
if w.shape != (h, CROSS_DIM):
raise RuntimeError("adapter %d shape %s != (%d,%d)" % (i, tuple(w.shape), h, CROSS_DIM))
return self
@torch.no_grad()
def ensure_on(self, device, dtype):
tag = (str(device), str(dtype))
if self._placed == tag:
return
self.to(device=device, dtype=dtype)
self._placed = tag
# ---------------------------------------------------------------------------
# Conditioning assembly (port of MSAdapter.generate, without diffusers pipeline)
# ---------------------------------------------------------------------------
class MSEmbeds:
"""Everything the attn2 patch needs, independent of CONDITIONING."""
def __init__(self, image_prompt_embeds, boxes, phrase_idxes, num_ref,
scale, mask_threshold, start_step, subject_scales=None):
self.image_prompt_embeds = image_prompt_embeds # [1, 4+16*num_ref, 2048] fp32
self.boxes = boxes # [1, num_ref, 4] fp32
self.phrase_idxes = phrase_idxes # [1, num_ref, 2] int
self.num_ref = num_ref
self.scale = float(scale)
self.mask_threshold = mask_threshold # None disables pseudo mask
self.start_step = int(start_step)
self.subject_scales = subject_scales
def minmax_normalize(batch_maps):
min_val = batch_maps.min(dim=-1, keepdim=True)[0].min(dim=-2, keepdim=True)[0]
max_val = batch_maps.max(dim=-1, keepdim=True)[0].max(dim=-2, keepdim=True)[0]
return (batch_maps - min_val) / (max_val - min_val + 1e-5)
def build_box_masks(boxes, seq_len_q, dtype, device):
"""Port of MaskedIPAttnProcessor2_0.prepare_attention_mask_qk (boxes path)."""
if boxes is None:
return None, None, None
batch_size, num_ref, _ = boxes.shape
nrow = int(seq_len_q ** 0.5)
box_idxes_start = torch.floor(boxes[:, :, 0:2] * nrow)
box_idxes_end = torch.ceil(boxes[:, :, 2:4] * nrow)
box_idxes = torch.cat([box_idxes_start, box_idxes_end], dim=-1)
box_masks = []
dummy_attention_mask = torch.ones((batch_size, seq_len_q), dtype=dtype, device=device)
for box_idx in box_idxes.unbind(dim=1):
xs, ys, xe, ye = box_idx.unbind(dim=1)
x_indices = torch.arange(nrow).unsqueeze(0).expand(batch_size, -1).to(device)
y_indices = torch.arange(nrow).unsqueeze(0).expand(batch_size, -1).to(device)
x_mask = ((x_indices >= xs.unsqueeze(1)) & (x_indices < xe.unsqueeze(1))).to(dtype)
y_mask = ((y_indices >= ys.unsqueeze(1)) & (y_indices < ye.unsqueeze(1))).to(dtype)
box_mask = torch.bmm(y_mask.unsqueeze(2), x_mask.unsqueeze(1)).reshape(batch_size, -1)
box_masks.append(box_mask)
dummy_attention_mask = torch.clamp(dummy_attention_mask - box_mask, min=0)
post_dummy = dummy_attention_mask.to(torch.bool)
image_mask = torch.stack(box_masks, dim=-1)
image_mask = image_mask.repeat_interleave(NUM_TOKENS, dim=-1)
image_mask = (1 - image_mask.to(dtype)) * -10000.0
dummy_mask = dummy_attention_mask.unsqueeze(-1).repeat_interleave(NUM_DUMMY_TOKENS, dim=-1)
dummy_mask = (1 - dummy_mask) * -10000.0
image_mask = torch.cat([dummy_mask, image_mask], dim=-1)
return image_mask, post_dummy, box_masks
def build_pseudo_masks(maps, mask_threshold, seq_len_q, dtype, device, num_ref):
"""Port of get_psuedo_attention_mask(). `maps` is a list of [B,num_ref,h,w]."""
text_maps = minmax_normalize(torch.stack(maps).mean(dim=0))
bsz, nr, h, w = text_maps.shape
seq_len_q = h * w
text_maps = text_maps.view(bsz, nr, -1).transpose(1, 2)
pseudo = (text_maps > mask_threshold).to(dtype)
pseudo_dummy = torch.ones((bsz, seq_len_q), dtype=dtype, device=device)
for i in range(nr):
pseudo_dummy = torch.clamp(pseudo_dummy - pseudo[..., i], min=0)
post_dummy = pseudo_dummy.to(torch.bool)
pseudo = pseudo.repeat_interleave(NUM_TOKENS, dim=-1)
pseudo = (1 - pseudo) * -10000.0
pseudo_dummy = pseudo_dummy.unsqueeze(-1).repeat_interleave(NUM_DUMMY_TOKENS, dim=-1)
pseudo_dummy = (1 - pseudo_dummy) * -10000.0
return torch.cat([pseudo_dummy, pseudo], dim=-1), post_dummy
# ---------------------------------------------------------------------------
# attn2 patch
# ---------------------------------------------------------------------------
# Module-level registry: ComfyUI's ModelPatcher.clone() deep-copies dict/list in
# model_options but passes arbitrary objects (functions, dicts of state) through
# unchanged. Closures capture only a string id, so cloned patchers share state.
_STATES = {}
def new_state(adapter, embeds):
sid = uuid.uuid4().hex
_STATES[sid] = {
"adapter": adapter,
"embeds": embeds,
"maps": {}, # per layer index -> list of text attention maps
"last_sigma": -1.0,
"calls": 0,
}
# keep the registry bounded across many node executions
if len(_STATES) > 64:
for k in list(_STATES)[:-32]:
_STATES.pop(k, None)
return sid
def _select_per_batch(t, batch, cond_or_uncond, zero_for_uncond=True):
"""Expand a [1,...] tensor to `batch` rows, zeroing uncond rows.
Upstream builds `cat([zeros_like(x), x])` for negative/positive. ComfyUI
instead reports the role of every batch chunk in transformer_options
["cond_or_uncond"] (0=cond, 1=uncond), so we build the same layout without
relying on a fixed chunk order.
"""
t = t.expand(batch, *t.shape[1:])
if not zero_for_uncond or cond_or_uncond is None:
return t
labels = list(cond_or_uncond)
if len(labels) == 0:
return t
if len(labels) != batch:
repeat = batch // len(labels)
if repeat * len(labels) != batch:
return t
labels = [x for x in labels for _ in range(repeat)]
keep = torch.tensor([1.0 if int(c) == 0 else 0.0 for c in labels],
device=t.device, dtype=t.dtype).view(batch, *([1] * (t.dim() - 1)))
return t * keep
def make_attn2_patch(state_id, layer_index, is_sentinel=False):
"""Return a ComfyUI `patches_replace["attn2"][(block, n, depth)]` callable.
Signature is (q, k, v, extra_options); q/k/v are already projected flat
tensors [B, L, inner_dim]. The enclosing BasicTransformerBlock applies
attn2.to_out and the residual, exactly like diffusers' processor.
"""
def attn2_patch(q, k, v, extra_options):
st = _STATES[state_id]
adapter = st["adapter"]
emb = st["embeds"]
# Detect a fresh sampling run by the (monotonically decreasing) sigma.
# All 70 patches see this; the first one resets, the rest no-op.
sigmas = extra_options.get("sigmas")
if sigmas is not None:
try:
sig = float(torch.as_tensor(sigmas).flatten()[0])
if sig > st["last_sigma"] + 1e-6:
st["maps"] = {}
st["calls"] = 0
st["last_sigma"] = sig
except Exception:
pass
if is_sentinel:
st["calls"] += 1
device, dtype = q.device, q.dtype
adapter.ensure_on(device, dtype)
bsz, seq_q, inner_dim = q.shape
heads = inner_dim // HEAD_DIM
head_dim = inner_dim // heads
scale_q = head_dim ** -0.5
use_pseudo = emb.mask_threshold is not None
def split_heads(t):
return t.view(t.shape[0], -1, heads, head_dim).transpose(1, 2)
qt = split_heads(q)
# ---- text branch -------------------------------------------------
kt, vt = split_heads(k), split_heads(v)
if not use_pseudo:
text_out = F.scaled_dot_product_attention(qt, kt, vt)
else:
scores = torch.matmul(qt * scale_q, kt.transpose(-1, -2))
text_probs = torch.softmax(scores.float(), dim=-1).to(dtype)
# upstream get_text_attention_maps() works on [B*heads, Lq, Lk]
flat = text_probs.reshape(bsz * heads, seq_q, -1).view(bsz, heads, seq_q, -1)
# Upstream builds cat([zeros, real]) for uncond/cond: uncond boxes
# and phrase indices are all zero. Reproduce that exactly.
boxes = _select_per_batch(emb.boxes.to(device=device, dtype=dtype), bsz,
extra_options.get("cond_or_uncond"))
pidx = _select_per_batch(emb.phrase_idxes.to(device=device, dtype=dtype), bsz,
extra_options.get("cond_or_uncond")).to(torch.int)
num_ref = boxes.shape[1]
h = w = int(seq_q ** 0.5)
maps = []
for i in range(bsz):
per_ref = []
for j in range(num_ref):
start_idx = int(pidx[i, j, 0].item())
end_idx = int(pidx[i, j, 1].item())
if start_idx == 0 and end_idx == 0:
per_ref.append(torch.zeros(seq_q, dtype=dtype, device=device))
else:
amap = flat[i, :, :, start_idx:end_idx]
per_ref.append(torch.mean(torch.mean(amap, dim=-1), dim=0))
maps.append(torch.stack(per_ref))
st["maps"].setdefault(layer_index, []).append(
torch.stack(maps).reshape(bsz, num_ref, h, w))
text_out = torch.matmul(text_probs, vt)
text_out = text_out.transpose(1, 2).reshape(bsz, seq_q, inner_dim).to(dtype)
# ---- image branch ------------------------------------------------
ip = emb.image_prompt_embeds.to(device=device, dtype=dtype)
ip = _select_per_batch(ip, bsz, extra_options.get("cond_or_uncond"))
ipk = split_heads(F.linear(ip, adapter.adapters[layer_index].to_k_ip.weight))
ipv = split_heads(F.linear(ip, adapter.adapters[layer_index].to_v_ip.weight))
boxes = _select_per_batch(emb.boxes.to(device=device, dtype=dtype), bsz,
extra_options.get("cond_or_uncond"))
rf_mask, dummy_mask, _ = build_box_masks(boxes, seq_q, dtype, device)
if use_pseudo:
maps_here = st["maps"].get(layer_index, [])
if len(maps_here) >= emb.start_step:
pm, pd = build_pseudo_masks(maps_here, emb.mask_threshold,
seq_q, dtype, device, boxes.shape[1])
rf_mask, dummy_mask = pm, pd
# [B, Lq, N] -> [B*heads, Lq, N]
if rf_mask is not None:
rf_mask = rf_mask.repeat_interleave(heads, dim=0).reshape(bsz * heads, seq_q, -1)
if dummy_mask is not None:
dummy_mask = dummy_mask.repeat_interleave(heads, dim=0)
qf = qt.reshape(bsz * heads, seq_q, head_dim)
kf = ipk.reshape(bsz * heads, -1, head_dim)
vf = ipv.reshape(bsz * heads, -1, head_dim)
scores = torch.bmm(qf * scale_q, kf.transpose(-1, -2))
if rf_mask is not None:
scores = scores + rf_mask
ip_probs = torch.softmax(scores.float(), dim=-1).to(dtype)
if dummy_mask is not None:
ip_probs = torch.where(dummy_mask.unsqueeze(-1), torch.zeros_like(ip_probs), ip_probs)
if emb.subject_scales is not None:
ss = torch.tensor(emb.subject_scales, dtype=dtype, device=device)
ss = ss.unsqueeze(0).unsqueeze(0).repeat_interleave(NUM_TOKENS, dim=-1)
dss = torch.ones((1, 1, 1), dtype=dtype, device=device).repeat_interleave(NUM_DUMMY_TOKENS, dim=-1)
ip_probs = ip_probs * torch.cat([dss, ss], dim=-1)
ip_out = torch.bmm(ip_probs, vf)
ip_out = ip_out.reshape(bsz, heads, seq_q, head_dim).transpose(1, 2).reshape(bsz, seq_q, inner_dim)
if emb.subject_scales is None:
return text_out + emb.scale * ip_out
return text_out + ip_out
return attn2_patch
# ---------------------------------------------------------------------------
# Weight -> ComfyUI block mapping
# ---------------------------------------------------------------------------
def enumerate_attn2(diffusion_model):
"""Return attn2 descriptors in ComfyUI *execution* order.
Each item: (key3, inner_dim, section) where key3 is the exact
(block_name, block_number, transformer_index) key ComfyUI looks up in
patches_replace["attn2"].
"""
import re
out = []
for name, module in diffusion_model.named_modules():
if not name.endswith(".attn2"):
continue
m = re.match(r"^(input_blocks|middle_block|output_blocks)\.(\d+)\..*?transformer_blocks\.(\d+)\.attn2$", name)
if m is None:
continue
section = {"input_blocks": "input", "middle_block": "middle", "output_blocks": "output"}[m.group(1)]
block_name = section
block_number = 0 if section == "middle" else int(m.group(2))
depth = int(m.group(3))
inner_dim = module.to_q.out_features
out.append(((block_name, block_number, depth), inner_dim, section))
# named_modules() already yields registration order == execution order
return out
def build_block_weights(attn2_desc, hidden_sizes):
"""Permute the 70 checkpoint groups into ComfyUI execution order.
Checkpoint order (diffusers registration): input | output | middle
ComfyUI execution order: input | middle | output
"""
n_input = sum(1 for x in attn2_desc if x[2] == "input")
n_middle = sum(1 for x in attn2_desc if x[2] == "middle")
n_output = sum(1 for x in attn2_desc if x[2] == "output")
total = n_input + n_middle + n_output
if total != len(hidden_sizes):
raise RuntimeError(
"MS-Diffusion: UNet has %d attn2 modules (%d input / %d middle / %d output) "
"but the checkpoint has %d adapter groups. This adapter only supports "
"stock SDXL (70 cross-attention layers)." % (total, n_input, n_middle, n_output, len(hidden_sizes)))
reordered = (list(range(0, n_input))
+ list(range(n_input + n_output, n_input + n_output + n_middle))
+ list(range(n_input, n_input + n_output)))
for pos, ckpt_idx in enumerate(reordered):
want = attn2_desc[pos][1]
got = hidden_sizes[ckpt_idx]
if want != got:
raise RuntimeError(
"MS-Diffusion: adapter mapping mismatch at attn2 #%d: UNet inner_dim=%d "
"but checkpoint group %d has hidden_size=%d. Refusing to run with wrong weights."
% (pos, want, ckpt_idx, got))
return reordered
def adapter_hidden_sizes(checkpoint_path):
"""Read just the 70 to_k_ip shapes so the loader node can build the module."""
sd = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
if not isinstance(sd, dict) or "ms_adapter" not in sd or "image_proj" not in sd:
raise RuntimeError(
"MS-Diffusion: %s is not an ms_adapter.bin checkpoint (needs "
"'image_proj', 'ms_adapter' and 'dummy_image_tokens' groups). "
"Pick ms_adapter.bin, not a generic IP-Adapter file." % checkpoint_path)
group = sd["ms_adapter"]
idx = sorted({int(k.split(".")[0]) for k in group})
try:
return [group["%d.to_k_ip.weight" % i].shape[0] for i in idx]
except KeyError as e:
raise RuntimeError("MS-Diffusion: unexpected ms_adapter key layout (%s). "
"Expected 70 to_k_ip/to_v_ip pairs." % e) from e
msdiffusion_native/nodes.py
+369
View File
@@ -0,0 +1,369 @@
"""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)",
}
pyproject.toml
+2 -1
View File
@@ -3,7 +3,8 @@ name = "comfyui_ms_diffusion"
description = "you can make story in comfyUI using MS-diffusion"
version = "1.1.0"
license = "LICENSE"
dependencies = ["tensorboard", "scipy", "opencv-python", "PyYAML"]
# Native chain depends only on torch, which ComfyUI already provides.
dependencies = []
[project.urls]
Repository = "https://github.com/smthemex/ComfyUI_MS_Diffusion"
requirements.txt
-3
View File
@@ -1,3 +0,0 @@
tensorboard
scipy
PyYAML