304 lines
12 KiB
Python
304 lines
12 KiB
Python
import dotenv
|
|
from os import environ as env
|
|
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-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 = option or {}
|
|
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 = option or {}
|
|
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 = option or {}
|
|
option["ignore_errors"] = ignore_errors
|
|
option["timeout"] = timeout_sec
|
|
option["retry"] = retry
|
|
return (option,)
|
|
|
|
|
|
class GenerateNAID:
|
|
def __init__(self):
|
|
dotenv.load_dotenv()
|
|
if "NAI_ACCESS_TOKEN" in env:
|
|
self.access_token = env["NAI_ACCESS_TOKEN"]
|
|
elif "NAI_ACCESS_KEY" in env:
|
|
print("ComfyUI_NAIDGenerator: NAI_ACCESS_KEY is deprecated. use NAI_ACCESS_TOKEN instead.")
|
|
access_key = env["NAI_ACCESS_KEY"]
|
|
elif "NAI_USERNAME" in env and "NAI_PASSWORD" in env:
|
|
print("ComfyUI_NAIDGenerator: NAI_USERNAME is deprecated. use NAI_ACCESS_TOKEN instead.")
|
|
username = env["NAI_USERNAME"]
|
|
password = env["NAI_PASSWORD"]
|
|
access_key = get_access_key(username, password)
|
|
else:
|
|
raise RuntimeError("Please ensure that NAI_API_TOKEN is set in ComfyUI/.env file.")
|
|
|
|
if not hasattr(self, "access_token"):
|
|
self.access_token = login(access_key)
|
|
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,
|
|
# "extra_noise_seed": seed, #TODO: find why it disappear
|
|
# "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:
|
|
params["reference_image_multiple"] = []
|
|
params["reference_information_extracted_multiple"] = []
|
|
params["reference_strength_multiple"] = []
|
|
|
|
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-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,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"GenerateNAID": GenerateNAID,
|
|
"ModelOptionNAID": ModelOption,
|
|
"Img2ImgOptionNAID": Img2ImgOption,
|
|
"InpaintingOptionNAID": InpaintingOption,
|
|
"VibeTransferOptionNAID": VibeTransferOption,
|
|
"NetworkOptionNAID": NetworkOption,
|
|
"MaskImageToNAID": ImageToNAIMask,
|
|
"PromptToNAID": PromptToNAID,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"GenerateNAID": "Generate ✒️🅝🅐🅘",
|
|
"ModelOptionNAID": "ModelOption ✒️🅝🅐🅘",
|
|
"Img2ImgOptionNAID": "Img2ImgOption ✒️🅝🅐🅘",
|
|
"InpaintingOptionNAID": "InpaintingOption ✒️🅝🅐🅘",
|
|
"VibeTransferOptionNAID": "VibeTransferOption ✒️🅝🅐🅘",
|
|
"NetworkOptionNAID": "NetworkOption ✒️🅝🅐🅘",
|
|
"MaskImageToNAID": "Convert Mask Image ✒️🅝🅐🅘",
|
|
"PromptToNAID": "Convert Prompt ✒️🅝🅐🅘",
|
|
}
|