457 lines
18 KiB
Python
457 lines
18 KiB
Python
import copy
|
|
import io
|
|
from pathlib import Path
|
|
import folder_paths
|
|
import zipfile
|
|
|
|
from .utils import *
|
|
|
|
class PromptToNAID:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return { "required": {
|
|
"text": ("STRING", { "forceInput":True, "multiline": True, "dynamicPrompts": False,}),
|
|
"weight_per_brace": ("FLOAT", { "default": 0.05, "min": 0.05, "max": 0.10, "step": 0.05 }),
|
|
}}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "convert"
|
|
CATEGORY = "NovelAI/utils"
|
|
def convert(self, text, weight_per_brace):
|
|
nai_prompt = prompt_to_nai(text, weight_per_brace)
|
|
return (nai_prompt,)
|
|
|
|
class ImageToNAIMask:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return { "required": { "image": ("IMAGE",) } }
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "convert"
|
|
CATEGORY = "NovelAI/utils"
|
|
def convert(self, image):
|
|
s = resize_to_naimask(image)
|
|
return (s,)
|
|
|
|
class ModelOption:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": (["safe-diffusion", "nai-diffusion", "nai-diffusion-furry", "nai-diffusion-2", "nai-diffusion-furry-3", "nai-diffusion-3"], { "default": "nai-diffusion-3" }),
|
|
},
|
|
"optional": { "option": ("NAID_OPTION",) },
|
|
}
|
|
|
|
RETURN_TYPES = ("NAID_OPTION",)
|
|
FUNCTION = "set_option"
|
|
CATEGORY = "NovelAI"
|
|
def set_option(self, model, option=None):
|
|
option = copy.deepcopy(option) if option else {}
|
|
option["model"] = model
|
|
return (option,)
|
|
|
|
class Img2ImgOption:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"strength": ("FLOAT", { "default": 0.70, "min": 0.01, "max": 0.99, "step": 0.01, "display": "number" }),
|
|
"noise": ("FLOAT", { "default": 0.00, "min": 0.00, "max": 0.99, "step": 0.02, "display": "number" }),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("NAID_OPTION",)
|
|
FUNCTION = "set_option"
|
|
CATEGORY = "NovelAI"
|
|
def set_option(self, image, strength, noise):
|
|
option = {}
|
|
option["img2img"] = (image, strength, noise)
|
|
return (option,)
|
|
|
|
class InpaintingOption:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"mask": ("IMAGE",),
|
|
"add_original_image": ("BOOLEAN", { "default": True }),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("NAID_OPTION",)
|
|
FUNCTION = "set_option"
|
|
CATEGORY = "NovelAI"
|
|
def set_option(self, image, mask, add_original_image):
|
|
option = {}
|
|
option["infill"] = (image, mask, add_original_image)
|
|
return (option,)
|
|
|
|
class VibeTransferOption:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"information_extracted": ("FLOAT", { "default": 1.0, "min": 0.01, "max": 1.0, "step": 0.01, "display": "number" }),
|
|
"strength": ("FLOAT", { "default": 0.6, "min": 0.01, "max": 1.0, "step": 0.01, "display": "number" }),
|
|
},
|
|
"optional": { "option": ("NAID_OPTION",) },
|
|
}
|
|
RETURN_TYPES = ("NAID_OPTION",)
|
|
FUNCTION = "set_option"
|
|
CATEGORY = "NovelAI"
|
|
def set_option(self, image, information_extracted, strength, option=None):
|
|
option = copy.deepcopy(option) if option else {}
|
|
if "vibe" not in option:
|
|
option["vibe"] = []
|
|
|
|
option["vibe"].append((image, information_extracted, strength))
|
|
return (option,)
|
|
|
|
class NetworkOption:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"ignore_errors": ("BOOLEAN", { "default": True }),
|
|
"timeout_sec": ("INT", { "default": 120, "min": 30, "max": 3000, "step": 1, "display": "number" }),
|
|
"retry": ("INT", { "default": 3, "min": 1, "max": 100, "step": 1, "display": "number" }),
|
|
},
|
|
"optional": { "option": ("NAID_OPTION",) },
|
|
}
|
|
RETURN_TYPES = ("NAID_OPTION",)
|
|
FUNCTION = "set_option"
|
|
CATEGORY = "NovelAI"
|
|
def set_option(self, ignore_errors, timeout_sec, retry, option=None):
|
|
option = copy.deepcopy(option) if option else {}
|
|
option["ignore_errors"] = ignore_errors
|
|
option["timeout"] = timeout_sec
|
|
option["retry"] = retry
|
|
return (option,)
|
|
|
|
|
|
class GenerateNAID:
|
|
def __init__(self):
|
|
self.access_token = get_access_token()
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"limit_opus_free": ("BOOLEAN", { "default": True }),
|
|
"width": ("INT", { "default": 832, "min": 64, "max": 1600, "step": 64, "display": "number" }),
|
|
"height": ("INT", { "default": 1216, "min": 64, "max": 1600, "step": 64, "display": "number" }),
|
|
"positive": ("STRING", { "default": ", best quality, amazing quality, very aesthetic, absurdres", "multiline": True, "dynamicPrompts": False }),
|
|
"negative": ("STRING", { "default": "lowres", "multiline": True, "dynamicPrompts": False }),
|
|
"steps": ("INT", { "default": 28, "min": 0, "max": 50, "step": 1, "display": "number" }),
|
|
"cfg": ("FLOAT", { "default": 5.0, "min": 0.0, "max": 10.0, "step": 0.1, "display": "number" }),
|
|
"smea": (["none", "SMEA", "SMEA+DYN"], { "default": "none" }),
|
|
"sampler": (["k_euler", "k_euler_ancestral", "k_dpmpp_2s_ancestral", "k_dpmpp_2m", "k_dpmpp_sde", "ddim"], { "default": "k_euler" }),
|
|
"scheduler": (["native", "karras", "exponential", "polyexponential"], { "default": "native" }),
|
|
"seed": ("INT", { "default": 0, "min": 0, "max": 9999999999, "step": 1, "display": "number" }),
|
|
"uncond_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 1.5, "step": 0.05, "display": "number" }),
|
|
"cfg_rescale": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.02, "display": "number" }),
|
|
},
|
|
"optional": { "option": ("NAID_OPTION",) },
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "NovelAI"
|
|
|
|
def generate(self, limit_opus_free, width, height, positive, negative, steps, cfg, smea, sampler, scheduler, seed, uncond_scale, cfg_rescale, option=None):
|
|
width, height = calculate_resolution(width*height, (width, height))
|
|
|
|
# ref. novelai_api.ImagePreset
|
|
params = {
|
|
"params_version": 1,
|
|
"width": width,
|
|
"height": height,
|
|
"scale": cfg,
|
|
"sampler": sampler,
|
|
"steps": steps,
|
|
"seed": seed,
|
|
"n_samples": 1,
|
|
"ucPreset": 3, #TODO: do I have to change it even if tags already typed by user?
|
|
"qualityToggle": False, #TODO: do I have to change it even if tags already typed by user?
|
|
"sm": (smea == "SMEA" or smea == "SMEA+DYN") and sampler != "ddim",
|
|
"sm_dyn": smea == "SMEA+DYN" and sampler != "ddim",
|
|
"dynamic_thresholding": False,
|
|
"controlnet_strength": 1.0,
|
|
"legacy": False,
|
|
"add_original_image": False,
|
|
"cfg_rescale": cfg_rescale,
|
|
"noise_schedule": scheduler,
|
|
"legacy_v3_extend": False, #TODO: find what it is
|
|
"uncond_scale": uncond_scale,
|
|
"negative_prompt": negative,
|
|
"reference_image_multiple": [], #NOTE: it is added on novelai webpage, even if ref.img not used
|
|
"reference_information_extracted_multiple": [],
|
|
"reference_strength_multiple": [],
|
|
"extra_noise_seed": seed, #NOTE: it uses for img2img but not sure okay to put on txt2img
|
|
# "decrisper": False,
|
|
}
|
|
model = "nai-diffusion-3"
|
|
action = "generate"
|
|
|
|
if option:
|
|
if "img2img" in option:
|
|
action = "img2img"
|
|
image, strength, noise = option["img2img"]
|
|
params["image"] = image_to_base64(resize_image(image, (width, height)))
|
|
params["strength"] = strength
|
|
params["noise"] = noise
|
|
elif "infill" in option:
|
|
action = "infill"
|
|
image, mask, add_original_image = option["infill"]
|
|
params["image"] = image_to_base64(resize_image(image, (width, height)))
|
|
params["mask"] = naimask_to_base64(resize_to_naimask(mask, (width, height)))
|
|
params["add_original_image"] = add_original_image
|
|
|
|
if "vibe" in option:
|
|
for vibe in option["vibe"]:
|
|
image, information_extracted, strength = vibe
|
|
params["reference_image_multiple"].append(image_to_base64(resize_image(image, (width, height))))
|
|
params["reference_information_extracted_multiple"].append(information_extracted)
|
|
params["reference_strength_multiple"].append(strength)
|
|
|
|
if "model" in option:
|
|
model = option["model"]
|
|
|
|
timeout = option["timeout"] if option and "timeout" in option else None
|
|
retry = option["retry"] if option and "retry" in option else None
|
|
|
|
if limit_opus_free:
|
|
pixel_limit = 1024*1024 if model in ("nai-diffusion-2", "nai-diffusion-furry-3", "nai-diffusion-3",) else 640*640
|
|
if width * height > pixel_limit:
|
|
max_width, max_height = calculate_resolution(pixel_limit, (width, height))
|
|
params["width"] = max_width
|
|
params["height"] = max_height
|
|
if steps > 28:
|
|
params["steps"] = 28
|
|
|
|
if sampler == "ddim" and model == "nai-diffusion-3":
|
|
params["sampler"] = "ddim_v3"
|
|
|
|
if action == "infill" and model != "nai-diffusion-2":
|
|
model = f"{model}-inpainting"
|
|
|
|
|
|
image = blank_image()
|
|
try:
|
|
zipped_bytes = generate_image(self.access_token, positive, model, action, params, timeout, retry)
|
|
zipped = zipfile.ZipFile(io.BytesIO(zipped_bytes))
|
|
image_bytes = zipped.read(zipped.infolist()[0]) # only support one n_samples
|
|
|
|
## save original png to comfy output dir
|
|
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path("NAI_autosave", self.output_dir)
|
|
file = f"{filename}_{counter:05}_.png"
|
|
d = Path(full_output_folder)
|
|
d.mkdir(exist_ok=True)
|
|
(d / file).write_bytes(image_bytes)
|
|
|
|
image = bytes_to_image(image_bytes)
|
|
except Exception as e:
|
|
if "ignore_errors" in option and option["ignore_errors"]:
|
|
print("ignore error:", e)
|
|
else:
|
|
raise e
|
|
|
|
return (image,)
|
|
|
|
|
|
def base_augment(access_token, output_dir, limit_opus_free, ignore_errors, req_type, image, options=None):
|
|
image = image.movedim(-1, 1)
|
|
w, h = (image.shape[3], image.shape[2])
|
|
image = image.movedim(1, -1)
|
|
|
|
if w * h > 1024 * 1024:
|
|
w, h = calculate_resolution(pixel_limit, (w, h))
|
|
base64_image = image_to_base64(resize_image(image, (w, h)))
|
|
result_image = blank_image()
|
|
try:
|
|
zipped_bytes = augment_image(access_token, req_type, w, h, base64_image, options=options)
|
|
zipped = zipfile.ZipFile(io.BytesIO(zipped_bytes))
|
|
image_bytes = zipped.read(zipped.infolist()[0]) # only support one n_samples
|
|
|
|
## save original png to comfy output dir
|
|
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path("NAI_autosave", output_dir)
|
|
file = f"{filename}_{counter:05}_.png"
|
|
d = Path(full_output_folder)
|
|
d.mkdir(exist_ok=True)
|
|
(d / file).write_bytes(image_bytes)
|
|
|
|
result_image = bytes_to_image(image_bytes)
|
|
except Exception as e:
|
|
if ignore_errors:
|
|
print("ignore error:", e)
|
|
else:
|
|
raise e
|
|
|
|
return (result_image,)
|
|
|
|
class RemoveBGAugment:
|
|
def __init__(self):
|
|
self.access_token = get_access_token()
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"limit_opus_free": ("BOOLEAN", { "default": True }),
|
|
"ignore_errors": ("BOOLEAN", { "default": False }),
|
|
},
|
|
}
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "augment"
|
|
CATEGORY = "NovelAI/director_tools"
|
|
def augment(self, image, limit_opus_free, ignore_errors):
|
|
return base_augment(self.access_token, self.output_dir, limit_opus_free, ignore_errors, "bg-removal", image)
|
|
|
|
class LineArtAugment:
|
|
def __init__(self):
|
|
self.access_token = get_access_token()
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"limit_opus_free": ("BOOLEAN", { "default": True }),
|
|
"ignore_errors": ("BOOLEAN", { "default": False }),
|
|
},
|
|
}
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "augment"
|
|
CATEGORY = "NovelAI/director_tools"
|
|
def augment(self, image, limit_opus_free, ignore_errors):
|
|
return base_augment(self.access_token, self.output_dir, limit_opus_free, ignore_errors, "lineart", image)
|
|
|
|
class SketchAugment:
|
|
def __init__(self):
|
|
self.access_token = get_access_token()
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"limit_opus_free": ("BOOLEAN", { "default": True }),
|
|
"ignore_errors": ("BOOLEAN", { "default": False }),
|
|
},
|
|
}
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "augment"
|
|
CATEGORY = "NovelAI/director_tools"
|
|
def augment(self, image, limit_opus_free, ignore_errors):
|
|
return base_augment(self.access_token, self.output_dir, limit_opus_free, ignore_errors, "sketch", image)
|
|
|
|
class ColorizeAugment:
|
|
def __init__(self):
|
|
self.access_token = get_access_token()
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"limit_opus_free": ("BOOLEAN", { "default": True }),
|
|
"ignore_errors": ("BOOLEAN", { "default": False }),
|
|
"defry": ("INT", { "default": 0, "min": 0, "max": 5, "step": 1, "display": "number" }),
|
|
"prompt": ("STRING", { "default": "", "multiline": True, "dynamicPrompts": False }),
|
|
},
|
|
}
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "augment"
|
|
CATEGORY = "NovelAI/director_tools"
|
|
def augment(self, image, limit_opus_free, ignore_errors, defry, prompt):
|
|
return base_augment(self.access_token, self.output_dir, limit_opus_free, ignore_errors, "colorize", image, options={ "defry": defry, "prompt": prompt })
|
|
|
|
class EmotionAugment:
|
|
def __init__(self):
|
|
self.access_token = get_access_token()
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
|
|
strength_list = ["normal", "slightly_weak", "weak", "even_weaker", "very_weak", "weakest"]
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"limit_opus_free": ("BOOLEAN", { "default": True }),
|
|
"ignore_errors": ("BOOLEAN", { "default": False }),
|
|
"mood": (["neutral", "happy", "sad", "angry", "scared",
|
|
"surprised", "tired", "excited", "nervous", "thinking",
|
|
"confused", "shy", "disgusted", "smug", "bored",
|
|
"laughing", "irritated", "aroused", "embarrassed", "worried",
|
|
"love", "determined", "hurt", "playful"], { "default": "neutral" }),
|
|
"strength": (s.strength_list, { "default": "normal" }),
|
|
"prompt": ("STRING", { "default": "", "multiline": True, "dynamicPrompts": False }),
|
|
},
|
|
}
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "augment"
|
|
CATEGORY = "NovelAI/director_tools"
|
|
def augment(self, image, limit_opus_free, ignore_errors, mood, strength, prompt):
|
|
prompt = f"{mood};;{prompt}"
|
|
defry = EmotionAugment.strength_list.index(strength)
|
|
return base_augment(self.access_token, self.output_dir, limit_opus_free, ignore_errors, "emotion", image, options={ "defry": defry, "prompt": prompt })
|
|
|
|
class DeclutterAugment:
|
|
def __init__(self):
|
|
self.access_token = get_access_token()
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"limit_opus_free": ("BOOLEAN", { "default": True }),
|
|
"ignore_errors": ("BOOLEAN", { "default": False }),
|
|
},
|
|
}
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "augment"
|
|
CATEGORY = "NovelAI/director_tools"
|
|
def augment(self, image, limit_opus_free, ignore_errors):
|
|
return base_augment(self.access_token, self.output_dir, limit_opus_free, ignore_errors, "declutter", image)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"GenerateNAID": GenerateNAID,
|
|
"ModelOptionNAID": ModelOption,
|
|
"Img2ImgOptionNAID": Img2ImgOption,
|
|
"InpaintingOptionNAID": InpaintingOption,
|
|
"VibeTransferOptionNAID": VibeTransferOption,
|
|
"NetworkOptionNAID": NetworkOption,
|
|
"MaskImageToNAID": ImageToNAIMask,
|
|
"PromptToNAID": PromptToNAID,
|
|
"RemoveBGNAID": RemoveBGAugment,
|
|
"LineArtNAID": LineArtAugment,
|
|
"SketchNAID": SketchAugment,
|
|
"ColorizeNAID": ColorizeAugment,
|
|
"EmotionNAID": EmotionAugment,
|
|
"DeclutterNAID": DeclutterAugment,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"GenerateNAID": "Generate ✒️🅝🅐🅘",
|
|
"ModelOptionNAID": "ModelOption ✒️🅝🅐🅘",
|
|
"Img2ImgOptionNAID": "Img2ImgOption ✒️🅝🅐🅘",
|
|
"InpaintingOptionNAID": "InpaintingOption ✒️🅝🅐🅘",
|
|
"VibeTransferOptionNAID": "VibeTransferOption ✒️🅝🅐🅘",
|
|
"NetworkOptionNAID": "NetworkOption ✒️🅝🅐🅘",
|
|
"MaskImageToNAID": "Convert Mask Image ✒️🅝🅐🅘",
|
|
"PromptToNAID": "Convert Prompt ✒️🅝🅐🅘",
|
|
"RemoveBGNAID": "Remove BG ✒️🅝🅐🅘",
|
|
"LineArtNAID": "LineArt ✒️🅝🅐🅘",
|
|
"SketchNAID": "SketchNAID ✒️🅝🅐🅘",
|
|
"ColorizeNAID": "Colorize ✒️🅝🅐🅘",
|
|
"EmotionNAID": "Emotion ✒️🅝🅐🅘",
|
|
"DeclutterNAID": "Declutter ✒️🅝🅐🅘",
|
|
}
|