501 lines
20 KiB
Python
501 lines
20 KiB
Python
import re
|
|
import numpy as np
|
|
import torch
|
|
from collections import defaultdict
|
|
import cv2
|
|
from PIL import Image, ImageOps, ImageFilter
|
|
from torchvision import transforms
|
|
from skimage.color import rgb2lab, deltaE_ciede2000
|
|
|
|
class prompt_sorting:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"string": ("STRING", {"multiline": True, "default": ""}),
|
|
"remove_duplicate_tags": ("BOOLEAN", {"default": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("string",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "advanced/AI_Assistant/text"
|
|
|
|
|
|
def execute(self,string,remove_duplicate_tags):
|
|
return (self.tidy_tags(string,remove_duplicate_tags),)
|
|
|
|
def tidy_tags(self,text, remove_duplicate_tags):
|
|
if not text or not isinstance(text,(str,list)):
|
|
return ""
|
|
|
|
# Convert list to string
|
|
if isinstance(text,list):
|
|
text = ",".join(map(str, text))
|
|
|
|
|
|
while True:
|
|
# Eliminate redundant commas and strip whitespace around commas
|
|
new_text = re.sub(r"[,\s\t\n]*,[,\s\t\n]*", ",", text)
|
|
|
|
# Fix double spaces
|
|
new_text = re.sub(r"\s\s+", " ", new_text)
|
|
if new_text == text:
|
|
break
|
|
text = new_text
|
|
|
|
text = text.strip(" ,\t\r\n")
|
|
|
|
# Remove duplicate tags, keeping the duplicate closest to the beginning of the prompt
|
|
if remove_duplicate_tags:
|
|
tags = text.split(',')
|
|
seen = set()
|
|
result = []
|
|
for tag in tags:
|
|
if tag not in seen:
|
|
seen.add(tag)
|
|
result.append(tag)
|
|
text = ",".join(result)
|
|
|
|
# Clean up BREAKs
|
|
# text = re.sub(r"[,\s\n]*BREAK[,\s\n]*"," BREAK ", text)
|
|
return text
|
|
|
|
|
|
class apply_lighting_effects:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"light_yaw": ("FLOAT", {
|
|
"default": 60,
|
|
"min": -180,
|
|
"max": 180,
|
|
"step": 0.1,
|
|
"display": "number"}),
|
|
"light_pitch": ("FLOAT", {
|
|
"default": -60,
|
|
"min": -90,
|
|
"max": 90,
|
|
"step": 0.1,
|
|
"display": "number"}),
|
|
"specular_power": ("FLOAT", {
|
|
"default": 30,
|
|
"min": 10,
|
|
"max": 100,
|
|
"step": 0.01,
|
|
"display": "number"}),
|
|
"normal_diffuse_strength": ("FLOAT", {
|
|
"default": 1,
|
|
"min": 0,
|
|
"max": 5,
|
|
"step": 0.01,
|
|
"display": "number"}),
|
|
"specular_highlights_strength": ("FLOAT", {
|
|
"default": 0.8,
|
|
"min": 0,
|
|
"max": 5,
|
|
"step": 0.01,
|
|
"display": "number"}),
|
|
"total_gain": ("FLOAT", {
|
|
"default": 0.6,
|
|
"min": 0,
|
|
"max": 1,
|
|
"step": 0.01,
|
|
"display": "number"}),
|
|
# "select_lighting_option": (
|
|
# [
|
|
# 'Directly Above',
|
|
# 'Directly Below',
|
|
# 'Directly Left',
|
|
# 'Directly Right',
|
|
# 'Upper Left Diagonal',
|
|
# 'Upper Right Diagonal'
|
|
# ], {
|
|
# "default": 'Directly Above'
|
|
# }),
|
|
},
|
|
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE","STRING", )
|
|
RETURN_NAMES = ("image","show_help",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "advanced/AI_Assistant/image"
|
|
DESCRIPTION = """
|
|
光源方向參考參數如下:
|
|
光源正前方=light_yaw:60 light_pitch:-60
|
|
光源左上方=light_yaw:40 light_pitch:-60
|
|
光源右上方=light_yaw:60 light_pitch:-40
|
|
光源左=light_yaw:0 light_pitch:0
|
|
光源右=light_yaw:90 light_pitch:0
|
|
光源下方=light_yaw:45 light_pitch:0
|
|
"""
|
|
|
|
|
|
def process(self, image, light_yaw, light_pitch, specular_power, normal_diffuse_strength,
|
|
specular_highlights_strength, total_gain):
|
|
show_help ="光源方向參考參數如下:\n光源正前方=light_yaw:60 light_pitch:-60\n光源左上方=light_yaw:40 light_pitch:-60\n光源右上方=light_yaw:60 light_pitch:-40\n光源左=light_yaw:0 light_pitch:0\n光源右=light_yaw:90 light_pitch:0\n光源下方=light_yaw:45 light_pitch:0\n"
|
|
|
|
return (self.lighting_effects(image, light_yaw, light_pitch, specular_power, normal_diffuse_strength,
|
|
specular_highlights_strength, total_gain),show_help,)
|
|
|
|
def lighting_effects(self, input_tensor, light_yaw, light_pitch, specular_power, normal_diffuse_strength,
|
|
specular_highlights_strength, total_gain):
|
|
|
|
# 元の次元数を記録する
|
|
original_dim = input_tensor.dim()
|
|
#print(original_dim,input_tensor.shape,input_tensor.shape[1])
|
|
|
|
# 入力テンソルが最低でも4次元であることを保証する（バッチ次元がなければ追加する）
|
|
if original_dim == 3:
|
|
input_tensor = input_tensor.unsqueeze(0) # バッチ次元を追加
|
|
|
|
# 轉換形狀為 (Batch, Channels, Height, Width)
|
|
input_tensor = input_tensor.permute(0, 3, 1, 2)
|
|
# 入力テンソルに3つ以上のチャンネルがある場合、追加のチャンネルをアルファと仮定し削除する
|
|
if input_tensor.shape[1] > 3:
|
|
input_tensor = input_tensor[:, :3, :, :] # 最初の3チャンネルのみを保持
|
|
|
|
# 入力テンソルを正規化して法線ベクトルを取得する
|
|
normal_tensor = torch.nn.functional.normalize(input_tensor, dim=1) # チャンネル次元に沿って正規化
|
|
|
|
# 光の方向ベクトルを計算する
|
|
light_direction = self.euler_to_vector(light_yaw, light_pitch, 0)
|
|
light_direction = light_direction.to(input_tensor.device) # 同じデバイスを保証する
|
|
|
|
# 光の方向とカメラの方向を入力テンソルの次元に合わせて拡張する
|
|
batch_size, _, height, width = input_tensor.shape
|
|
light_direction = light_direction.view(1, 3, 1, 1).expand(batch_size, -1, height, width)
|
|
camera_direction = torch.tensor([0, 0, 1], dtype=torch.float32, device=input_tensor.device)
|
|
camera_direction = camera_direction.view(1, 3, 1, 1).expand(batch_size, -1, height, width)
|
|
|
|
# 拡散成分を計算する
|
|
diffuse = torch.sum(normal_tensor * light_direction, dim=1, keepdim=True)
|
|
diffuse = torch.clamp(diffuse, 0, 1)
|
|
|
|
# 鏡面成分を計算する
|
|
half_vector = torch.nn.functional.normalize(light_direction + camera_direction, dim=1)
|
|
specular = torch.sum(normal_tensor * half_vector, dim=1, keepdim=True)
|
|
specular = torch.pow(torch.clamp(specular, 0, 1), specular_power)
|
|
|
|
# 拡散成分と鏡面成分を組み合わせて、強度とゲインを適用する
|
|
output_tensor = (diffuse * normal_diffuse_strength + specular * specular_highlights_strength) * total_gain
|
|
output_tensor = output_tensor.squeeze(1) # keepdim=Trueで追加されたチャンネル次元を削除
|
|
|
|
# 初めに追加されたバッチ次元があれば削除する
|
|
if original_dim == 3:
|
|
output_tensor = output_tensor.squeeze(0)
|
|
|
|
return output_tensor
|
|
|
|
def euler_to_vector(self, yaw, pitch, roll):
|
|
# オイラー角から方向ベクトルを計算
|
|
yaw_rad = np.radians(yaw)
|
|
pitch_rad = np.radians(pitch)
|
|
roll_rad = np.radians(roll)
|
|
|
|
cos_pitch = np.cos(pitch_rad)
|
|
sin_pitch = np.sin(pitch_rad)
|
|
cos_yaw = np.cos(yaw_rad)
|
|
sin_yaw = np.sin(yaw_rad)
|
|
|
|
direction = np.array([
|
|
sin_yaw * cos_pitch,
|
|
-sin_pitch,
|
|
cos_yaw * cos_pitch
|
|
])
|
|
return torch.from_numpy(direction).float()
|
|
|
|
class clean_prompt_tags:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"string": ("STRING", {"multiline": True, "default": "", "forceInput": True}),
|
|
"Clean_Prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING","STRING",)
|
|
RETURN_NAMES = ("string","show_help")
|
|
FUNCTION = "execute"
|
|
CATEGORY = "advanced/AI_Assistant/text"
|
|
DESCRIPTION = """
|
|
粗略的移除工具:
|
|
|
|
EX:
|
|
base_prompt = masterpiece, best quality, monochrome, greyscale,blue_hair, blue_sky, blue_dress, pink_eyes, pink_bow, pink_panties
|
|
|
|
removeprompt = pink, red, orange, brown, yellow, green, blue, purple, blonde, colored skin, white hair,
|
|
|
|
輸出則為:masterpiece, best quality, monochrome, greyscale
|
|
"""
|
|
|
|
def execute(self,string,Clean_Prompt):
|
|
show_help ="粗略的移除工具:\n\nEX:\n\nbase_prompt = masterpiece, best quality, monochrome, greyscale,blue_hair, blue_sky, blue_dress, pink_eyes, pink_bow, pink_panties\n\nremoveprompt = pink, red, orange, brown, yellow, green, blue, purple, blonde, colored skin, white hair,\n\n輸出則為:\nmasterpiece, best quality, monochrome, greyscale"
|
|
return (self.remove_prompt(string,Clean_Prompt),show_help)
|
|
|
|
def remove_prompt(self,base_prompt, remove_prompt,):
|
|
|
|
# 移除末尾多餘的逗號並使用正規表達式來分割標籤
|
|
remove_prompt = re.sub(r',\s*$', '', remove_prompt).replace('\n', '')
|
|
|
|
# 使用正規表達式來分割標籤，忽略逗號後的多餘空格
|
|
prompt_list = re.split(r',\s*', base_prompt)
|
|
color_list = re.split(r',\s*', remove_prompt)
|
|
|
|
# カラータグを除去します。
|
|
cleaned_tags = [tag for tag in prompt_list if all(color.lower() not in tag.lower() for color in color_list)]
|
|
|
|
return ", ".join(cleaned_tags)
|
|
|
|
class prompt_blacklist:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"string": ("STRING", {"multiline": True, "default": "", "forceInput": True}),
|
|
"Blacklist_Prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING","STRING",)
|
|
RETURN_NAMES = ("string","show_help")
|
|
FUNCTION = "execute"
|
|
CATEGORY = "advanced/AI_Assistant/text"
|
|
DESCRIPTION = """
|
|
黑名單功能，分別輸入需移除的提示以逗號分隔即可
|
|
|
|
EX:
|
|
base_prompt= masterpiece, best quality,lineart,sketch,transparent background
|
|
Blacklist Prompt=lineart,sketch,
|
|
輸出則為:masterpiece, best quality, transparent background
|
|
"""
|
|
def execute(self,Blacklist_Prompt,string):
|
|
show_help ="黑名單功能，分別輸入需移除的提示以逗號分隔即可\n\nEX:\nbase_prompt= masterpiece, best quality,lineart,sketch,transparent background\nBlacklist Prompt=lineart,sketch,\n輸出則為:\n masterpiece, best quality, transparent background"
|
|
return (self.execute_prompt(Blacklist_Prompt,string),show_help)
|
|
|
|
def execute_prompt(self,execute_tags, base_prompt):
|
|
prompt_list = re.split(r',\s*', base_prompt)
|
|
# execute_tagsを除去
|
|
filtered_tags = [tag for tag in prompt_list if tag not in execute_tags]
|
|
# 最終的なプロンプトを生成
|
|
return ", ".join(filtered_tags)
|
|
|
|
class resize_image_sdxl_ratio:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("INT","INT","STRING",)
|
|
RETURN_NAMES = ("Target Width","Target Height","show_help")
|
|
FUNCTION = "execute"
|
|
CATEGORY = "advanced/AI_Assistant/image"
|
|
DESCRIPTION = """
|
|
自動輸出對應的SDXL目標尺寸與ImageScale(圖像縮放)搭配使用並開啟剪裁功能
|
|
"""
|
|
def execute(self,image):
|
|
width,height=self.resize_image_aspect_ratio(image)
|
|
show_help ="自動輸出對應的SDXL目標尺寸與ImageScale(圖像縮放)搭配使用並開啟剪裁功能"
|
|
return (width, height, show_help)
|
|
|
|
def ResolutionFromImage(self, image):
|
|
_, H, W, _ = image.shape
|
|
return (W, H)
|
|
|
|
def resize_image_aspect_ratio(self,image):
|
|
# 元の画像サイズを取得
|
|
original_width, original_height = self.ResolutionFromImage(image)
|
|
|
|
# アスペクト比を計算
|
|
aspect_ratio = original_width / original_height
|
|
|
|
# 標準のアスペクト比サイズを定義
|
|
sizes = {
|
|
1: (1024, 1024), # 正方形
|
|
4/3: (1152, 896), # 横長画像
|
|
3/2: (1216, 832),
|
|
16/9: (1344, 768),
|
|
21/9: (1568, 672),
|
|
3/1: (1728, 576),
|
|
1/4: (512, 2048), # 縦長画像
|
|
1/3: (576, 1728),
|
|
9/16: (768, 1344),
|
|
2/3: (832, 1216),
|
|
3/4: (896, 1152)
|
|
}
|
|
|
|
# 最も近いアスペクト比を見つける
|
|
closest_aspect_ratio = min(sizes.keys(), key=lambda x: abs(x - aspect_ratio))
|
|
target_width, target_height = sizes[closest_aspect_ratio]
|
|
|
|
return (target_width, target_height)
|
|
|
|
class noline_process:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
},
|
|
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE","STRING", )
|
|
RETURN_NAMES = ("image","show_help",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "advanced/AI_Assistant/image"
|
|
DESCRIPTION = """
|
|
將輸入圖像移除線條並將顏色取近似值用於I2I
|
|
"""
|
|
|
|
def process(self, image,):
|
|
show_help ="將輸入圖像移除線條並將顏色取近似值用於I2I\n"
|
|
|
|
return (self.noline_process(image),show_help,)
|
|
|
|
def noline_process(self, input_image):
|
|
|
|
#ToPILImage
|
|
input_image = input_image.permute(0, 3, 1, 2)
|
|
input_image = input_image.squeeze(0)
|
|
input_image = transforms.ToPILImage()(input_image)
|
|
|
|
def get_major_colors(image, threshold_percentage):
|
|
if image.mode != 'RGB':
|
|
image = image.convert('RGB')
|
|
color_count = defaultdict(int)
|
|
for pixel in image.getdata():
|
|
color_count[pixel] += 1
|
|
total_pixels = image.width * image.height
|
|
return [(color, count) for color, count in color_count.items() if (count / total_pixels) >= threshold_percentage]
|
|
|
|
def consolidate_colors(major_colors, threshold):
|
|
colors_lab = [rgb2lab(np.array([[color]], dtype=np.float32)/255.0).reshape(3) for color, _ in major_colors]
|
|
i = 0
|
|
while i < len(colors_lab):
|
|
j = i + 1
|
|
while j < len(colors_lab):
|
|
if deltaE_ciede2000(colors_lab[i], colors_lab[j]) < threshold:
|
|
if major_colors[i][1] >= major_colors[j][1]:
|
|
major_colors[i] = (major_colors[i][0], major_colors[i][1] + major_colors[j][1])
|
|
major_colors.pop(j)
|
|
colors_lab.pop(j)
|
|
else:
|
|
major_colors[j] = (major_colors[j][0], major_colors[j][1] + major_colors[i][1])
|
|
major_colors.pop(i)
|
|
colors_lab.pop(i)
|
|
continue
|
|
j += 1
|
|
i += 1
|
|
return major_colors
|
|
|
|
def generate_distant_colors(consolidated_colors, distance_threshold):
|
|
consolidated_lab = [rgb2lab(np.array([color], dtype=np.float32) / 255.0).reshape(3) for color, _ in consolidated_colors]
|
|
max_attempts = 10000
|
|
for _ in range(max_attempts):
|
|
random_rgb = np.random.randint(0, 256, size=3)
|
|
random_lab = rgb2lab(np.array([random_rgb], dtype=np.float32) / 255.0).reshape(3)
|
|
if all(deltaE_ciede2000(base_color_lab, random_lab) > distance_threshold for base_color_lab in consolidated_lab):
|
|
return tuple(random_rgb)
|
|
return (128, 128, 128)
|
|
|
|
def line_color(image, mask, new_color):
|
|
data = np.array(image)
|
|
data[mask, :3] = new_color
|
|
return Image.fromarray(data)
|
|
|
|
def replace_color(image, color_1, blur_radius=2):
|
|
data = np.array(image)
|
|
original_shape = data.shape
|
|
channels = original_shape[2] if len(original_shape) > 2 else 1
|
|
data = data.reshape(-1, channels)
|
|
color_1 = np.array(color_1)
|
|
matches = np.all(data[:, :3] == color_1, axis=1)
|
|
mask = np.zeros(data.shape[0], dtype=bool)
|
|
|
|
while np.any(matches):
|
|
new_matches = np.zeros_like(matches)
|
|
for i in range(len(data)):
|
|
if matches[i]:
|
|
x, y = divmod(i, original_shape[1])
|
|
neighbors = [(x, y-1), (x, y+1), (x-1, y), (x+1, y)]
|
|
valid_neighbors = [data[nx * original_shape[1] + ny, :3] for nx, ny in neighbors if 0 <= nx < original_shape[0] and 0 <= ny < original_shape[1] and not matches[nx * original_shape[1] + ny]]
|
|
if valid_neighbors:
|
|
new_color = np.mean(valid_neighbors, axis=0).astype(np.uint8)
|
|
data[i, :3] = new_color
|
|
mask[i] = True
|
|
else:
|
|
new_matches[i] = True
|
|
matches = new_matches
|
|
if not np.any(matches):
|
|
break
|
|
|
|
data = data.reshape(original_shape)
|
|
mask = mask.reshape(original_shape[:2])
|
|
result_image = Image.fromarray(data, 'RGBA')
|
|
blurred_image = result_image.filter(ImageFilter.GaussianBlur(radius=blur_radius))
|
|
blurred_data = np.array(blurred_image)
|
|
np.copyto(data, blurred_data, where=mask[..., None])
|
|
return Image.fromarray(data, 'RGBA')
|
|
|
|
def DoG_filter(image, kernel_size=0, sigma=1.0, k_sigma=2.0, gamma=1.5):
|
|
g1 = cv2.GaussianBlur(image, (kernel_size, kernel_size), sigma)
|
|
g2 = cv2.GaussianBlur(image, (kernel_size, kernel_size), sigma * k_sigma)
|
|
return g1 - gamma * g2
|
|
|
|
def XDoG_filter(image, kernel_size=0, sigma=1.4, k_sigma=1.6, epsilon=0, phi=10, gamma=0.98):
|
|
epsilon /= 255
|
|
dog = DoG_filter(image, kernel_size, sigma, k_sigma, gamma)
|
|
dog /= dog.max()
|
|
e = 1 + np.tanh(phi * (dog - epsilon))
|
|
e[e >= 1] = 1
|
|
return (e * 255).astype('uint8')
|
|
|
|
def binarize_image(image):
|
|
_, binarized = cv2.threshold(image, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
|
|
return binarized
|
|
|
|
|
|
def process_XDoG(image):
|
|
kernel_size = 0
|
|
sigma = 1.4
|
|
k_sigma = 1.6
|
|
epsilon = 0
|
|
phi = 10
|
|
gamma = 0.98
|
|
|
|
# 將 PIL Image 對象轉換為灰度模式的 numpy 數組
|
|
image_array = np.array(image.convert("L"))
|
|
xdog_image = XDoG_filter(image_array, kernel_size, sigma, k_sigma, epsilon, phi, gamma)
|
|
binarized_image = binarize_image(xdog_image)
|
|
final_image = Image.fromarray(binarized_image)
|
|
return final_image
|
|
|
|
rgb_image = input_image.convert("RGBA")
|
|
lineart = process_XDoG(input_image).convert('L')
|
|
lineart = lineart.point(lambda x: 0 if x < 200 else 255)
|
|
lineart = ImageOps.invert(lineart)
|
|
kernel = np.ones((3, 3), np.uint8)
|
|
lineart = cv2.dilate(np.array(lineart), kernel, iterations=1)
|
|
lineart = Image.fromarray(lineart)
|
|
mask = np.array(lineart) == 255
|
|
major_colors = get_major_colors(rgb_image, threshold_percentage=0.05)
|
|
major_colors = consolidate_colors(major_colors, 10)
|
|
new_color_1 = generate_distant_colors(major_colors, 100)
|
|
filled_image = line_color(rgb_image, mask, new_color_1)
|
|
replace_color_image = replace_color(filled_image, new_color_1, 2).convert('RGB')
|
|
#ToTensor
|
|
replace_color_image=transforms.ToTensor()(replace_color_image)
|
|
replace_color_image = replace_color_image.unsqueeze(0)
|
|
replace_color_image = replace_color_image.permute(0, 2, 3, 1)
|
|
#print(replace_color_image.shape)
|
|
return replace_color_image
|