Files
smthemex-ComfyUI_MS_Diffusion/MS_Diffusion.py
T
2024-07-03 12:00:20 +08:00

946 lines
49 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# !/usr/bin/env python
# -*- coding: UTF-8 -*-
import numpy as np
import torch
import os
import random
import sys
import re
from PIL import ImageFont
import yaml
from PIL import Image
from huggingface_hub import hf_hub_download
from diffusers import (StableDiffusionXLPipeline, DiffusionPipeline, DDIMScheduler, ControlNetModel,
KDPM2AncestralDiscreteScheduler, LMSDiscreteScheduler,
AutoPipelineForInpainting, DPMSolverMultistepScheduler, DPMSolverSinglestepScheduler,
EulerDiscreteScheduler, HeunDiscreteScheduler, UNet2DConditionModel,
AutoPipelineForText2Image, StableDiffusionXLControlNetImg2ImgPipeline, 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 nodes import LoadImage
from .utils.utils import get_comic
from .utils.style_template import styles
STYLE_NAMES = list(styles.keys())
import diffusers
LoadImage=LoadImage()
dif_version = str(diffusers.__version__)
dif_version_int = int(dif_version.split(".")[1])
device = (
"cuda"
if torch.cuda.is_available()
else "mps" if torch.backends.mps.is_available() else "cpu"
)
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)
fonts_path = os.path.join(dir_path, "fonts")
fonts_lists = os.listdir(fonts_path)
control_paths = []
for search_path in folder_paths.get_folder_paths("diffusers"):
if os.path.exists(search_path):
for root, subdir, files in os.walk(search_path, followlinks=True):
if "config.json" in files:
control_paths.append(os.path.relpath(root, start=search_path))
control_paths = ([z for z in control_paths if "controlnet-canny-sdxl-1.0" in z]
+ [p for p in control_paths if "MistoLine" in p]
+ [o for o in control_paths if "lcm-sdxl" in o]
+ [Q for Q in control_paths if "controlnet-openpose-sdxl-1.0" in Q]
+ [Z for Z in control_paths if "controlnet-scribble-sdxl-1.0" in Z]
+ [a for a in control_paths if "controlnet-depth-sdxl-1.0" in a]
+[b for b in control_paths if "controlnet-tile-sdxl-1.0" in b]
+[c for c in control_paths if "controlnet-zoe-depth-sdxl-1.0" in c]
+[d for d in control_paths if "sdxl-controlnet-seg " in d])
if control_paths:
control_paths = ["none"] + [x for x in control_paths if x]
else:
control_paths = ["none",]
diff_paths = []
for search_path in folder_paths.get_folder_paths("diffusers"):
if os.path.exists(search_path):
for root, subdir, files in os.walk(search_path, followlinks=True):
if "model_index.json" in files:
diff_paths.append(os.path.relpath(root, start=search_path))
if diff_paths:
diff_paths = ["none"] + [x for x in diff_paths if x]
else:
diff_paths = ["none",]
clip_paths = []
for search_path in folder_paths.get_folder_paths("clip_vision"):
if os.path.exists(search_path):
for root, subdir, files in os.walk(search_path, followlinks=True):
if "config.json" in files:
clip_paths.append(os.path.relpath(root, start=search_path))
if diff_paths:
clip_paths = ["none"] + [x for x in clip_paths if x]
else:
clip_paths = ["none",]
def pil2tensor(img):
img_convert_to_numpy = np.array(img) # (32, 32, 3)
tensor = torch.tensor(img_convert_to_numpy.transpose(2, 0, 1) / 255) # torch.Size([3, 32, 32])
return tensor
def apply_style_positive(style_name: str, positive: str):
p, n = styles.get(style_name, styles[style_name])
#print(p, "test0", n)
return p.replace("{prompt}", positive)
def apply_style(style_name: str, positives: list, negative: str = ""):
p, n = styles.get(style_name, styles[style_name])
#print(p,"test1",n)
return [
p.replace("{prompt}", positive) for positive in positives
], n + " " + negative
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 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 remove_punctuation_from_strings(lst):
pattern = r"[\W]+$" # 匹配字符串末尾的所有非单词字符
return [re.sub(pattern, '', s) for s in lst]
def format_punctuation_from_strings(lst):
pattern = r"[\W]+$" # 匹配字符串末尾的所有非单词字符
return [re.sub(pattern, ';', s) for s in lst]
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 setup_seed(seed):
torch.manual_seed(seed)
if device == "cuda":
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
random.seed(seed)
torch.backends.cudnn.deterministic = True
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(img, width, height):
samples = img.movedim(-1, 1)
img = common_upscale(samples, width, height, "nearest-exact", "center")
samples = img.movedim(1, -1)
img = tensor_to_image(samples)
return img
def input_size_adaptation_output(img_tensor,base_in, width, height):
#basein=512
if width == height:
img_pil = nomarl_upscale(img_tensor, base_in, base_in) # 2pil
else:
if min(1,width/ height)==1: #高
r=height/base_in
img_pil = nomarl_upscale(img_tensor, round(width/r), base_in) # 2pil
else: #宽
r=width/base_in
img_pil = nomarl_upscale(img_tensor, base_in, round(height/r)) # 2pil
img_pil.resize((width,height))
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):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"ckpt_name": (["none"]+folder_paths.get_filename_list("checkpoints"),),
"diffuser_model": (diff_paths,),
"repo_id": ("STRING", {"default": "stabilityai/stable-diffusion-xl-base-1.0"}),
"clip_vision_local":(clip_paths,),
"clip_vision_repo": ("STRING", {"default": "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"}),
"controlnet_model_path": (control_paths,),
"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 = ("MODEL","MODEL","MODEL","MODEL", "STRING",)
RETURN_NAMES = ("pipe","ms_model","image_encoder", "image_processor","info",)
FUNCTION = "ms_model_loader"
CATEGORY = "MSdiffusion"
def ms_model_loader(self, ckpt_name,diffuser_model,repo_id,clip_vision_local, clip_vision_repo, controlnet_model_path,lora, lora_scale, trigger_words,scheduler,
):
repo_id = instance_path(diffuser_model, repo_id,"diffusers")
clip_vision=instance_path(clip_vision_local, clip_vision_repo,"clip_vision")
if clip_vision=="none":
raise "need clip_vision"
scheduler_choice = get_scheduler(scheduler)
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
original_config_file = os.path.join(dir_path, 'config', 'sd_xl_base.yaml')
original_config_file = get_instance_path(original_config_file)
if controlnet_model_path != "none": #启用control
controlnet_model_path = get_instance_path(get_local_path(file_path, controlnet_model_path,"diffusers"))
controlnet = ControlNetModel.from_pretrained(controlnet_model_path, variant="fp16",
use_safetensors=True,
torch_dtype=torch.float16).to(device)
if repo_id == "none":
if ckpt_name == "none":
raise "need choice checkpoints or fill repo or choice diffuser model"
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
repo_id = get_instance_path(ckpt_path)
if dif_version_int >= 28:
pipe = StableDiffusionXLControlNetPipeline.from_single_file(
repo_id, original_config=original_config_file, controlnet=controlnet,
torch_dtype=torch.float16,
add_watermarker=False, ).to(device)
else:
pipe = StableDiffusionXLControlNetPipeline.from_single_file(
repo_id, original_config_file=original_config_file, controlnet=controlnet,
torch_dtype=torch.float16,
add_watermarker=False,
).to(device)
else:
pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
repo_id, torch_dtype=torch.float16, add_watermarker=False, controlnet=controlnet,
).to(device)
else: # 非controlnet模式
if repo_id == "none":
if ckpt_name == "none":
raise "need choice checkpoints or fill repo or choice diffuser model"
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
repo_id = get_instance_path(ckpt_path)
if dif_version_int >= 28:
pipe = StableDiffusionXLPipeline.from_single_file(
repo_id, original_config=original_config_file, torch_dtype=torch.float16,
add_watermarker=False, ).to(device)
else:
pipe = StableDiffusionXLPipeline.from_single_file(
repo_id, original_config_file=original_config_file, torch_dtype=torch.float16,
add_watermarker=False,
).to(device)
else:
pipe = StableDiffusionXLPipeline.from_pretrained(
repo_id, torch_dtype=torch.float16, add_watermarker=False,
).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_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(dir_path, "weights")
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()
#print(type(image_processor))
image_encoder = CLIPVisionModelWithProjection.from_pretrained(clip_vision).to(device, dtype=torch.float16)
# print(type(image_encoder))
image_encoder_projection_dim = image_encoder.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.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)
info = str(";".join(
[lora, trigger_words,controlnet_model_path,image_encoder_type,image_proj_type]))
return (pipe,ms_model,image_encoder,image_processor, info,)
class MSdiffusion_Sampler:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"info": ("STRING", {"forceInput": True, "default": ""}),
"pipe": ("MODEL",),
"ms_model": ("MODEL",),
"image_encoder": ("MODEL",),
"image_processor": ("MODEL",),
"character_prompt": ("STRING", {"multiline": True,
"default": "[woman] wearing a white T-shirt, blue loose hair.\n"
"[man] wearing a suit,black hair."}),
"scene_prompts": ("STRING", {"multiline": True,
"default": "[woman] and [man] have breakfast,medium shot;\n[woman] go to company;\n[man] have lunch,medium shot;\n[NC] a living room."}),
"split_prompt": ("STRING", {"default": ""}),
"negative_prompt": ("STRING", {"multiline": True,
"default": "bad anatomy, bad hands, missing fingers, extra fingers, "
"three hands, three legs, bad arms, missing legs, "
"missing arms, poorly drawn face, bad face, fused face, "
"cloned face, three crus, fused feet, fused thigh, "
"extra crus, ugly fingers, horn,"
"amputation, disconnected limbs"}),
"img_style": (
["No_style", "Realistic", "Japanese_Anime", "Digital_Oil_Painting", "Pixar_Disney_Character",
"Photographic", "Comic_book",
"Line_art", "Black_and_White_Film_Noir", "Isometric_Rooms"],),
"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}),
"role_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"}),
"layout_guidance":("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",),}
}
RETURN_TYPES = ("IMAGE", "STRING",)
RETURN_NAMES = ("image", "prompt_array",)
FUNCTION = "ms_sampler"
CATEGORY = "MSdiffusion"
def main_normal(self,prompts_dual,pipe,phrases,ms_model,input_images,num_samples,steps,seed,negative_prompt,role_scale,image_encoder,cfg,image_processor,
boxes,mask_threshold,start_step,image_proj_type,image_encoder_type,drop_grounding_tokens,height,width,img_style):
image_ouput = []
for i, prompt in enumerate(prompts_dual):
prompt = apply_style_positive(img_style, prompt)
#print(prompt)
# used to get the attention map, return zero if the phrase is not in the prompt
phrase_idxes = [self.get_phrases_idx(pipe.tokenizer, phrases[0], prompt)]
eot_idxes = [[get_eot_idx(pipe.tokenizer, prompt)] * len(phrases[0])]
# print(phrase_idxes, eot_idxes)
images = ms_model.generate(pipe=pipe, pil_images=[input_images], num_samples=num_samples,
num_inference_steps=steps,
seed=seed,
prompt=[prompt], negative_prompt=negative_prompt, scale=role_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_ouput += images
return image_ouput
def main_control(self,prompts_dual,control_img_list,width,height,pipe,phrases,ms_model,input_images,num_samples,steps,seed,negative_prompt,role_scale,image_encoder,cfg,
image_processor,boxes,mask_threshold,start_step,image_proj_type,image_encoder_type,drop_grounding_tokens,controlnet_scale,img_style):
j = 0
image_ouput=[]
for i, prompt in enumerate(prompts_dual):
control_image = control_img_list[j]
j += 1
prompt=apply_style_positive(img_style,prompt)
#print(prompt)
# used to get the attention map, return zero if the phrase is not in the prompt
phrase_idxes = [self.get_phrases_idx(pipe.tokenizer, phrases[0], prompt)]
eot_idxes = [[get_eot_idx(pipe.tokenizer, prompt)] * len(phrases[0])]
# print(phrase_idxes, eot_idxes)
images = ms_model.generate(pipe=pipe, pil_images=[input_images], num_samples=num_samples,
num_inference_steps=steps,
seed=seed,
prompt=[prompt], negative_prompt=negative_prompt, scale=role_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, ontrolnet_conditioning_scale=controlnet_scale)
image_ouput += images
return image_ouput
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 add_lora(self,trigger_words,lora,prompts_list):
# 添加Lora trigger
add_trigger_words = "," + trigger_words + " " + "style" + ";"
prompts_list = remove_punctuation_from_strings(prompts_list)
if lora not in lora_lightning_list: # 加速lora不需要trigger
prompts_list = [item + add_trigger_words for item in prompts_list]
return prompts_list
def ms_sampler(self,image,info,pipe,ms_model,image_encoder,image_processor, character_prompt, scene_prompts,split_prompt,negative_prompt,img_style,seed, steps,
cfg,role_scale, mask_threshold, start_step,controlnet_scale,width,height,layout_guidance,guidance_list,**kwargs):
lora, trigger_words,controlnet_model_path,image_encoder_type,image_proj_type = info.split(";")
guidance_list=guidance_list.strip().split(";")
box_a = guidance_list[0].split(",")
box_b = guidance_list[1].split(",")
box_a_float=[float(x) for x in box_a]
box_b_float = [float(x) for x in box_b]
box_add=[box_a_float,box_b_float]
#print(box_a_float, box_b_float,box_add)
#处理prompt
if split_prompt:
scene_prompts.replace("\n", "").replace(split_prompt, ";\n").strip()
character_prompt.replace("\n", "").replace(split_prompt, ";\n").strip()
else:
scene_prompts.strip()
character_prompt.strip()
if "\n" not in scene_prompts:
scene_prompts.replace(";", ";\n").strip()
if "\n" in character_prompt:
if character_prompt.count("\n") > 1:
character_prompt.replace("\n", "").replace("[", "\n[").strip()
if character_prompt.count("\n") > 1:
character_prompt.replace("\n", "").replace("[", "\n[", 2).strip() # 多行角色在这里强行转为双角色
# 从角色列表获取角色方括号信息
char_origin_f = character_prompt.splitlines()
char_describe = [char.replace("]", " ").replace("[", " ") for char in char_origin_f] #del character's []
char_origin = [char.split("]")[0] + "]" for char in char_origin_f] # [[role1],[role2]]
role_str=[x.replace("]", "").replace("[", "") for x in char_origin] # [role1,role2]
prompts_origin = scene_prompts.splitlines()
control_check=prompts_origin
prompts_origin_unuse, negative_prompt = apply_style(img_style, prompts_origin, negative_prompt) #获取 style p prompt
num_samples = 1 #批次
#image = kwargs["image"]
d1, _, _, _ = image.size()
if d1 == 1:
prompt_single_NC = [prompt for prompt in prompts_origin if "[NC]" in prompt] # NC场景列表,单人,
prompt_single_NC=[prompt.replace("[NC]","") for prompt in prompt_single_NC ] # 去除符号
index_single_NC = [index for index, prompt in enumerate(prompts_origin) if "[NC]" in prompt] # NC场景index
prompts_origin = [prompt for prompt in prompts_origin if role_str[0] in prompt] # 排除 nc
input_images = [nomarl_upscale(image, width, height)]
torch.cuda.empty_cache()
# 加入角色描述
prompt_c = char_origin_f[0].replace("]", " ").replace("[", " ")
role = char_origin[0].replace("]", "").replace("[", "")
prompts_dual = [x.replace(char_origin[0],prompt_c) for x in prompts_origin]
prompts_dual = [x.replace("]", "").replace("[", "") for x in prompts_dual]
if layout_guidance:
boxes=random.choice(box_add)
boxes=[[boxes]] #随机选1个自定义的
#boxes = [[[0.25, 0.25, 0.75, 0.75]]] # 1 role
else:
boxes = [[[0., 0., 0., 0.]]] # 1 role
phrases = [[role]]
drop_grounding_tokens = [0] # set to 1 if you want to drop the grounding tokens
if len(char_origin_f) != 1: # 单人
raise "need one role prompts "
if controlnet_model_path != "none":
control_image = kwargs["control_image"]
d1, _, _, _ = control_image.size()
if d1!=len(control_check):
raise "The number of scene prompts that must match the control_image"
control_img = list(torch.chunk(control_image, chunks=d1)) #tensor
control_img = [input_size_adaptation_output(control_image, 768, width, height) for control_image in control_img ] #pil
control_img_list=[img for index, img in enumerate(control_img) if index not in index_single_NC ] #del NC
if lora != "none":
prompts_dual=self.add_lora(trigger_words, lora, prompts_dual)
image_main=self.main_control(prompts_dual,control_img_list,width,height,pipe,phrases,ms_model,input_images,num_samples,steps,seed,negative_prompt,role_scale,image_encoder,cfg,
image_processor,boxes,mask_threshold,start_step,image_proj_type,image_encoder_type,drop_grounding_tokens,controlnet_scale,img_style)
if prompt_single_NC:
if lora != "none":
prompt_single_NC = self.add_lora(trigger_words, lora, prompt_single_NC)
control_img_list_NC= [img for index, img in enumerate(control_img) if index in index_single_NC ] # NC
input_images= [Image.new('RGB', (width, height), (255, 255, 255))]
image_nc = self.main_control(prompt_single_NC, control_img_list_NC, width, height, pipe, phrases,
ms_model, input_images, num_samples, steps, seed, negative_prompt,
role_scale, image_encoder, cfg,
image_processor, boxes, mask_threshold, start_step, image_proj_type,
image_encoder_type, drop_grounding_tokens, controlnet_scale,img_style)
jj = 0
for i in index_single_NC: # 重新将NC场景插入原序列
img = image_nc[jj]
image_main.insert(int(i), img)
jj += 1
image_list = narry_list(image_main)
else:
image_list = narry_list(image_main)
else: #no controlnet
if lora != "none":
prompts_dual=self.add_lora(trigger_words, lora, prompts_dual)
image_main=self.main_normal( prompts_dual, pipe, phrases, ms_model, input_images, num_samples, steps, seed,
negative_prompt, role_scale, image_encoder, cfg, image_processor,
boxes, mask_threshold, start_step, image_proj_type, image_encoder_type,
drop_grounding_tokens, height, width,img_style)
if prompt_single_NC:
if lora != "none":
prompt_single_NC = self.add_lora(trigger_words, lora, prompt_single_NC)
input_images= [Image.new('RGB', (width, height), (255, 255, 255))]
image_nc = self.main_normal( prompt_single_NC, pipe, phrases, ms_model, input_images, num_samples, steps, seed,
negative_prompt, role_scale, image_encoder, cfg, image_processor,
boxes, mask_threshold, start_step, image_proj_type, image_encoder_type,
drop_grounding_tokens, height, width,img_style)
jj = 0
for i in index_single_NC: # 重新将NC场景插入原序列
img = image_nc[jj]
image_main.insert(int(i), img)
jj += 1
image_list = narry_list(image_main)
else:
prompt_dual_NC = [prompt for prompt in prompts_origin if
"[NC]" in prompt] # NC场景列表,双人,
prompt_dual_NC=[prompt.replace("[NC]","") for prompt in prompt_dual_NC] # 去除符号
index_dual_NC = [index for index, prompt in enumerate(prompts_origin) if
"[NC]" in prompt] # NC场景index
prompts_origin_role_a = [prompt for prompt in prompts_origin if
char_origin[0] in prompt and char_origin[1] not in prompt] # 排除 nc,仅单人A
prompts_origin_role_a_index = [index for index, prompt in enumerate( prompts_origin) if
char_origin[0] in prompt and char_origin[1] not in prompt]
prompts_origin_role_b = [prompt for prompt in prompts_origin if
char_origin[1] in prompt and char_origin[0] not in prompt] # 排除 nc,仅单人B
prompts_origin_role_b_index = [index for index, prompt in enumerate(prompts_origin) if
char_origin[1] in prompt and char_origin[0] not in prompt]
prompts_origin = [prompt for prompt in prompts_origin if
(char_origin[0] in prompt and char_origin[1] in prompt)] # 排除单人,nc,仅双人
#print(prompt_dual_NC, "prompt_dual_NC",index_dual_NC, "index_dual_NC",prompts_origin_role_a ,"prompts_origin_role_a ",prompts_origin_role_a_index,"prompts_origin_role_a_index",prompts_origin_role_b,"prompts_origin_role_b",prompts_origin,"prompts_origin")
img_list = list(torch.chunk(image, chunks=d1))
input_images = [nomarl_upscale(img, width, height) for img in img_list]
input_images_a =[input_images[0]]
input_images_b = [input_images[1]]
# get role name and describe
role_a = role_str[0]
char_describe_a= char_describe[0]
role_b =role_str[1]
char_describe_b = char_describe[1]
# 双角色场景列表 prompt加描述
prompts_dual = [item.replace(char_origin[0],char_describe_a).replace(char_origin[1],char_describe_b) for item in prompts_origin]
torch.cuda.empty_cache()
if layout_guidance:
boxes=[box_add]
#boxes = [[[0., 0.25, 0.4, 0.75], [0.6, 0.25, 1., 0.75]]] # man+women
else:
boxes = [[[0., 0., 0., 0.], [0., 0., 0., 0.]]] # used if you want no layout guidance
phrases = [[role_a, role_b]]
drop_grounding_tokens = [0] # set to 1 if you want to drop the grounding tokens
image_new=Image.new('RGB', (width, height), (255, 255, 255))
if controlnet_model_path != "none":
control_image = kwargs["control_image"]
d1, _, _, _ = control_image.size()
if d1 != len(control_check):
raise "The number of scene prompts that must match the control_image"
control_img = list(torch.chunk(control_image, chunks=d1)) # tensor
control_img = [input_size_adaptation_output(control_image, 768, width, height) for control_image in
control_img] # pil
control_img_list = [img for index, img in enumerate(control_img) if
index not in index_dual_NC] # del NC
if lora != "none":
prompts_dual=self.add_lora(trigger_words, lora, prompts_dual)
image_main=self.main_control(prompts_dual,control_img_list,width,height,pipe,phrases,ms_model,input_images,num_samples,steps,seed,negative_prompt,role_scale,image_encoder,cfg,
image_processor,boxes,mask_threshold,start_step,image_proj_type,image_encoder_type,drop_grounding_tokens,controlnet_scale,img_style)
if prompt_dual_NC:
if lora != "none":
prompt_dual_NC = self.add_lora(trigger_words, lora, prompt_dual_NC)
control_img_list_NC = [img for index, img in enumerate(control_img) if
index in index_dual_NC] # NC
input_images = [image_new,image_new]
image_nc = self.main_control(prompt_dual_NC, control_img_list_NC, width, height, pipe, phrases,
ms_model, input_images, num_samples, steps, seed, negative_prompt,
role_scale, image_encoder, cfg,
image_processor, boxes, mask_threshold, start_step, image_proj_type,
image_encoder_type, drop_grounding_tokens, controlnet_scale,img_style)
jj = 0
for i in index_dual_NC: # 重新将NC场景插入原序列
img = image_nc[jj]
image_main.insert(int(i), img)
jj += 1
if prompts_origin_role_a:
if lora != "none":
prompts_origin_role_a = self.add_lora(trigger_words, lora, prompts_origin_role_a)
prompts_origin_role_a =[x.replace(char_origin[0],char_describe_a) for x in prompts_origin_role_a] # 加入角色描述
control_img_list = [img for index, img in enumerate(control_img) if
index in prompts_origin_role_a_index] # role a
if layout_guidance:
boxes = random.choice(box_add)
boxes = [[boxes]] # 随机选1个自定义的
#boxes = [[[0.25, 0.25, 0.75, 0.75]]] # 1 role
else:
boxes = [[[0., 0., 0., 0.]]] # 1 role
phrases = [[role_a]]
input_images = input_images_a
image_dual_a = self.main_control(prompts_origin_role_a, control_img_list, width, height, pipe, phrases,
ms_model, input_images, num_samples, steps, seed, negative_prompt,
role_scale, image_encoder, cfg,
image_processor, boxes, mask_threshold, start_step, image_proj_type,
image_encoder_type, drop_grounding_tokens, controlnet_scale,img_style)
a = 0
for i in prompts_origin_role_a_index: # 重新将NC场景插入原序列
img = image_dual_a[a]
image_main.insert(int(i), img)
a += 1
if prompts_origin_role_b:
if lora != "none":
prompts_origin_role_b = self.add_lora(trigger_words, lora, prompts_origin_role_b)
prompts_origin_role_b = [x.replace(char_origin[1], char_describe_b) for x in
prompts_origin_role_b] # 加入角色描述
control_img_list = [img for index, img in enumerate(control_img) if
index in prompts_origin_role_b_index] # role b
if layout_guidance:
boxes = random.choice(box_add)
boxes = [[boxes]] # 随机选1个自定义的
#boxes = [[[0.25, 0.25, 0.75, 0.75]]] # 1 role
else:
boxes = [[[0., 0., 0., 0.]]] # 1 role
phrases = [[role_b]]
input_images = input_images_b
image_dual_b = self.main_control(prompts_origin_role_b, control_img_list, width, height, pipe, phrases,
ms_model, input_images, num_samples, steps, seed, negative_prompt,
role_scale, image_encoder, cfg,
image_processor, boxes, mask_threshold, start_step, image_proj_type,
image_encoder_type, drop_grounding_tokens, controlnet_scale,img_style)
b = 0
for i in prompts_origin_role_b_index: # 重新将NC场景插入原序列
img = image_dual_b[b]
image_main.insert(int(i), img)
b += 1
image_list = narry_list(image_main)
else:
if lora != "none":
prompts_dual = self.add_lora(trigger_words, lora, prompts_dual)
image_main = self.main_normal(prompts_dual, pipe, phrases, ms_model, input_images, num_samples, steps,
seed,
negative_prompt, role_scale, image_encoder, cfg, image_processor,
boxes, mask_threshold, start_step, image_proj_type, image_encoder_type,
drop_grounding_tokens, height, width,img_style)
if prompt_dual_NC:
if lora != "none":
prompt_dual_NC = self.add_lora(trigger_words, lora, prompt_dual_NC)
input_images = [image_new,image_new]
image_nc = self.main_normal( prompt_dual_NC, pipe, phrases, ms_model, input_images, num_samples,
steps, seed,
negative_prompt, role_scale, image_encoder, cfg, image_processor,
boxes, mask_threshold, start_step, image_proj_type, image_encoder_type,
drop_grounding_tokens, height, width,img_style)
jj = 0
for i in index_dual_NC: # 重新将NC场景插入原序列
img = image_nc[jj]
image_main.insert(int(i), img)
jj += 1
if prompts_origin_role_a:
if lora != "none":
prompts_origin_role_a = self.add_lora(trigger_words, lora, prompts_origin_role_a)
prompts_origin_role_a = [x.replace(char_origin[0], char_describe_a) for x in
prompts_origin_role_a] # 加入角色描述
if layout_guidance:
boxes = random.choice(box_add)
boxes = [[boxes]] # 随机选1个自定义的
#boxes = [[[0.25, 0.25, 0.75, 0.75]]] # 1 role
else:
boxes = [[[0., 0., 0., 0.]]] # 1 role
phrases = [[role_a]]
input_images = input_images_a
image_role_a = self.main_normal(prompts_origin_role_a, pipe, phrases, ms_model, input_images, num_samples,
steps, seed,
negative_prompt, role_scale, image_encoder, cfg, image_processor,
boxes, mask_threshold, start_step, image_proj_type, image_encoder_type,
drop_grounding_tokens, height, width,img_style)
jj = 0
for i in prompts_origin_role_a_index: # 重新将NC场景插入原序列
img = image_role_a[jj]
image_main.insert(int(i), img)
jj += 1
if prompts_origin_role_b:
if lora != "none":
prompts_origin_role_b = self.add_lora(trigger_words, lora, prompts_origin_role_b)
prompts_origin_role_b = [x.replace(char_origin[1], char_describe_b) for x in
prompts_origin_role_b] # 加入角色描述
if layout_guidance:
boxes = random.choice(box_add)
boxes = [[boxes]] # 随机选1个自定义的
#boxes = [[[0.25, 0.25, 0.75, 0.75]]] # 1 role
else:
boxes = [[[0., 0., 0., 0.]]] # 1 role
phrases = [[role_b]]
input_images = input_images_b
image_role_b = self.main_normal(prompts_origin_role_b, pipe, phrases, ms_model, input_images, num_samples,
steps, seed,
negative_prompt, role_scale, image_encoder, cfg, image_processor,
boxes, mask_threshold, start_step, image_proj_type, image_encoder_type,
drop_grounding_tokens, height, width,img_style)
jj = 0
for i in prompts_origin_role_b_index: # 重新将NC场景插入原序列
img = image_role_b[jj]
image_main.insert(int(i), img)
jj += 1
image_list = narry_list(image_main)
image = torch.from_numpy(np.fromiter(image_list, np.dtype((np.float32, (height, width, 3)))))
torch.cuda.empty_cache()
return (image, scene_prompts,)
class MS_Comic_Type:
@classmethod
def INPUT_TYPES(s):
return {"required": {"image": ("IMAGE",),
"scene_prompts": ("STRING", {"multiline": True, "forceInput": True, "default": ""}),
"fonts_list": (fonts_lists,),
"text_size": ("INT", {"default": 40, "min": 1, "max": 100}),
"comic_type": (["Four_Pannel", "Classic_Comic_Style"],),
"split_lines": ("STRING", {"default": ";"}),
}}
RETURN_TYPES = ("IMAGE",)
ETURN_NAMES = ("image",)
FUNCTION = "comic_gen"
CATEGORY = "MSdiffusion"
def comic_gen(self, image, scene_prompts, fonts_list, text_size, comic_type, split_lines):
result = [item for index, item in enumerate(image)]
total_results = narry_list_pil(result)
font_choice = os.path.join(dir_path, "fonts", fonts_list)
captions = scene_prompts.splitlines()
if len(captions) > 1:
# captions = [caption.replace("(", "").replace(")", "") if "(" or ")" in caption else caption
# for caption in captions] # del ()
captions = [caption.replace("[NC]", "") for caption in captions]
captions = [caption.replace("]", "").replace("[", "") for caption in captions]
# captions = [
# caption.split("#")[-1] if "#" in caption else caption
# for caption in captions
# ]
else:#trans
prompt_array = scene_prompts.replace(split_lines, "\n")
captions = prompt_array.splitlines()
font = ImageFont.truetype(font_choice, text_size)
images = (
get_comic(total_results, comic_type, captions=captions, font=font)
+ total_results
)
images = phi2narry(images[0])
return (images,)
NODE_CLASS_MAPPINGS = {
"MSdiffusion_Model_Loader": MSdiffusion_Model_Loader,
"MSdiffusion_Sampler": MSdiffusion_Sampler,
"MS_Comic_Type":MS_Comic_Type
}
NODE_DISPLAY_NAME_MAPPINGS = {
"MSdiffusion_Model_Loader": "MSdiffusion_Model_Loader",
"MSdiffusion_Sampler": "MSdiffusion_Sampler",
"MS_Comic_Type":"MS_Comic_Type"
}