nodes/bedovyy-ComfyUI_NAIDGenerator
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #45 from Mooshieblob1/master

Browse Source 
Added Character Reference
This commit is contained in:
Bedovyy
2025-10-12 06:48:05 +09:00
committed by GitHub
parent f40dc475a9 52cdcfb53f
commit 55d886809b
2 changed files with 363 additions and 332 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+66 -101
View File
@@ -1,6 +1,6 @@
# ComfyUI_NAIDGenerator
A [ComfyUI](https://github.com/comfyanonymous/ComfyUI) extension for generating image via NovelAI API.
A [ComfyUI](https://github.com/comfyanonymous/ComfyUI) extension for generating images via the NovelAI API.
## Installation
@@ -9,89 +9,116 @@ A [ComfyUI](https://github.com/comfyanonymous/ComfyUI) extension for generating
## Setting up NAI account
Before using the nodes, you should set NAI_ACCESS_TOKEN on `ComfyUI/.env` file.
Before using the nodes, you should set `NAI_ACCESS_TOKEN` in a `.env` file located in your main `ComfyUI` directory.
`ComfyUI/.env`
```
NAI_ACCESS_TOKEN=<ACCESS_TOKEN>
NAI_ACCESS_TOKEN=<YOUR_ACCESS_TOKEN>
```
You can get persistent API token by **User Settings > Account > Get Persistent API Token** on NovelAI webpage.
You can get a persistent API token by navigating to **User Settings > Account > Get Persistent API Token** on the NovelAI website.
Otherwise, you can get access token which is valid for 30 days using [novelai-api](https://github.com/Aedial/novelai-api).
Otherwise, you can get an access token which is valid for 30 days using [novelai-api](https://github.com/Aedial/novelai-api).
## Usage
The nodes are located at `NovelAI` category.
The nodes are located in the `NovelAI` category.
![image](https://github.com/bedovyy/ComfyUI_NAIDGenerator/assets/137917911/8ab1ecc0-2ba8-4e38-8810-727e50a20923)
### Txt2img
Simply connect `GenerateNAID` node and `SaveImage` node.
Simply connect the `GenerateNAID` node to a `SaveImage` node.
![generate](https://github.com/bedovyy/ComfyUI_NAIDGenerator/assets/137917911/1328896d-7d4b-4d47-8ec2-d1c4e8e2561c)
Note that all generated images via `GeneratedNAID` node are saved as `output/NAI_autosave_12345_.png` for keeping original metadata.
**Note:** All generated images via the `GenerateNAID` node are automatically saved to `output/NAI_autosave/NAI_autosave_#####_.png` to preserve their original metadata.
### Img2img
Connect `Img2ImgOptionNAID` node to `GenerateNAID` node and put original image.
Connect an `Img2ImgOptionNAID` node to the `option` input of the `GenerateNAID` node and provide a source image.
![image](https://github.com/bedovyy/ComfyUI_NAIDGenerator/assets/137917911/15ff8961-4f6b-4f23-86bf-34b86ace45c0)
Note that width and height of the source image will be resized to generation size.
**Note:** The width and height of the source image will be resized to the generation size.
### Inpainting
Connect `InpaintingOptionNAID` node to `GenerateNAID` node and put original image and mask image.
Connect an `InpaintingOptionNAID` node to the `GenerateNAID` node and provide a source image and a mask.
![image](https://github.com/bedovyy/ComfyUI_NAIDGenerator/assets/137917911/5ed1ad77-b90e-46be-8c37-9a5ee0935a3d)
Note that both source image and mask will be resized fit to generation size.
(You don't need `MaskImageToNAID` node to convert mask image to NAID mask image.)
**Note:** Both the source image and mask will be automatically resized to fit the generation size.
### Vibe Transfer
Connect `VibeTransferOptionNAID` node to `GenerateNAID` node and put reference image.
Connect a `VibeTransferOptionNAID` node to the `GenerateNAID` node and provide a reference image to transfer its style and feel.
![Comfy_workflow](https://github.com/bedovyy/ComfyUI_NAIDGenerator/assets/137917911/8c6c1c2e-f29d-42a1-b615-439155cb3164)
You can also relay Img2ImgOption on it.
You can also chain it with other options, like Img2Img.
![image](https://github.com/bedovyy/ComfyUI_NAIDGenerator/assets/137917911/acf0496c-8c7c-48f4-9530-18e6a23669d5)
Note that width and height of the source images will be resized to generation size. **This will change aspect ratio of source images.**
#### Multiple Vibe Transfer
Just connect multiple `VibeTransferOptionNAID` nodes to `GenerateNAID` node.
Connect multiple `VibeTransferOptionNAID` nodes to combine their influences.
![preview_vibe_2](https://github.com/user-attachments/assets/2d56c0f7-bcd5-48ff-b436-012ea43604fe)
### Character Reference
Use the `CharacterReferenceOptionNAID` node to guide the generation using a single reference image for character identity and/or style.
![image](https://github.com/user-attachments/assets/9f21eff1-163e-49f4-9e4a-85ea25b29988)
- **style_aware:** If enabled, it attempts to copy both the character's features and the artistic style.
- **fidelity:** Controls how strictly the generation should adhere to the reference image. The developers state that `primary_strength` is always kept at `1.0`, while `secondary_strength` is calculated as $1.0 - \text{fidelity}$.
**Note:** The reference image will be automatically letterboxed to an accepted NAI canvas size to preserve its aspect ratio.
### ModelOption
The default model of `GenerateNAID` node is `nai-diffusion-3`(NAI Diffusion Anime V3).
The default model of the `GenerateNAID` node is `nai-diffusion-4-5-full`. To change the model, connect a `ModelOptionNAID` node.
If you want to change model, put `ModelOptionNAID` node to `GenerateNAID` node.
Available V4+ models include:
- `nai-diffusion-4-curated-preview`
- `nai-diffusion-4-full`
- `nai-diffusion-4-5-curated`
- `nai-diffusion-4-5-full`
![ModelOption](https://github.com/bedovyy/ComfyUI_NAIDGenerator/assets/137917911/0b484edb-bcb5-428a-b2af-1372a9d7a34f)
### NetworkOption
You can set timeout or retry option from `NetworkOption` node.
Moreover, you can ignore error by `ignore_errors`. In that case, the result will be 1x1 size grayscale image.
Without this node, the request never retry and wait response forever, and stop the queue when error occurs
You can set timeout and retry options using the `NetworkOption` node. You can also set `ignore_errors` to prevent the queue from stopping on an API error; if an error occurs, a blank 1x1 image will be output.
![preview_network](https://github.com/user-attachments/assets/d82b0ff2-c57c-4870-9024-8d78261a8fea)
**Note that if you set timeout too short, you may not get image but spend Anlas.**
**Note:** If you set the timeout too short, you may not receive an image but could still be charged Anlas.
### Anlas Tracker
This extension now includes Anlas tracking to monitor your usage.
**1. Console Output (Automatic)**
All generation and director tool nodes will automatically print your Anlas balance before and after the operation in the console where you launched ComfyUI.
```
[NovelAI] Anlas (pre-gen): 10000
[NovelAI] Generation cost: 20 Anlas
[NovelAI] Anlas (post-gen): 9980
```
**2. Visual Node**
A new node, **`Anlas Tracker ✒️🅝🅐🅘`**, is available in the `NovelAI/utils` category. You can use it to display your current Anlas balance directly in your workflow.
Connect its `anlas_string` output to a display node (e.g., "Show Text" from the [WAS Node Suite](https://github.com/WASasquatch/was-node-suite-comfyui)) to see the value. Use the `trigger` input to control when the balance is checked.
![image](https://github.com/user-attachments/assets/387799cd-36bd-4b5d-9746-bfd74db09f74) <!-- Placeholder for actual image -->
### PromptToNAID
ComfyUI use `()` or `(word:weight)` for emphasis, but NovelAI use `{}` and `[]`. This node convert ComfyUI's prompt to NovelAI's.
Optionally, you can choose weight per brace. If you set `weight_per_brace` to 0.10, `(word:1.1)` will convert to `{word}` instead of `{{word}}`.
ComfyUI uses `()` or `(word:weight)` for emphasis, while NovelAI uses `{}` and `[]`. This node, found in `NovelAI/utils`, converts ComfyUI's prompt syntax to NovelAI's.
![image](https://github.com/bedovyy/ComfyUI_NAIDGenerator/assets/137917911/25c48350-7268-4d6f-81fe-9eb080fc6e5a)
@@ -99,91 +126,29 @@ Optionally, you can choose weight per brace. If you set `weight_per_brace` to 0.
![image](https://github.com/user-attachments/assets/e205a51e-59dc-4d5a-94c8-29715ed98739)
You can find director tools like `LineArtNAID` or `EmotionNAID` on NovelAI > director_tools.
You can find director tools like `LineArtNAID`, `EmotionNAID`, and `RemoveBGNAID` in the `NovelAI/director_tools` category.
![augment_example](https://github.com/user-attachments/assets/5833e9fb-f92e-4d53-9069-58ca8503a3e7)
### V4 Support (Preview)
### V4 / V4.5 Support
The node now supports NAI's V4 architecture through the nai-diffusion-4-curated-preview model. This is a preview release of V4 with some limitations:
- **Important Notes:**
- This is a preview version of V4 and some features are limited
- Inpainting will automatically use V3 model (but works with V4-generated images)
- Vibe transfer is not yet supported with V4 preview (will be available with full V4 release)
- Full V4 feature support will come with the official V4 release
### V4.5 Support (Curated Preview)
Support has been added for **NAI Diffusion 4.5 Curated Preview**, an updated version of V4 with further improvements in detail, contrast, and prompt responsiveness.
- **Model Name:**
```python
model = "nai-diffusion-4-5-curated-preview"
```
- **Availability:**
Selectable through `ModelOptionNAID` node under the name **NAI Diffusion 4.5 Curated Preview**.
- **Compatibility Notes:**
- Works the same as V4 preview, with the same limitations:
- Inpainting will still default to V4 backend
- Vibe transfer is not yet supported
- Prompt formatting remains the same as for V4 (`V4BasePrompt` and `V4NegativePrompt` nodes are compatible)
#### New Model Option
NAI Diffusion V4 Curated Preview is now available in the ModelOptionNAID node:
```python
model = "nai-diffusion-4-curated-preview"
```
The nodes fully support NAI's V4 and V4.5 model architecture.
#### V4 Prompt Handling
Two new nodes have been added for V4 prompt handling:
Two new nodes have been added in `NovelAI/v4` for V4/V4.5 prompt handling:
##### V4BasePrompt
- **`V4BasePrompt`**: Handles the positive prompt.
- **`V4NegativePrompt`**: Handles the negative prompt.
A node for handling V4 positive prompts:
#### Example V4 / V4.5 Workflow
Here's a basic setup for a V4/V4.5 model:
```
V4BasePrompt -----> positive
GenerateNAID
```
##### V4NegativePrompt
A node for handling V4 negative prompts:
```
GenerateNAID
V4NegativePrompt -> negative
GenerateNAID
```
#### Example V4 Workflow
Here's a basic V4 setup:
```
V4BasePrompt -----> positive
V4NegativePrompt -> negative GenerateNAID
ModelOption ------> option
```
#### Work In Progress Features
The following V4 features are currently in development:
```python
"""
- V4PromptConfig: Advanced prompt configuration
- Coordinate-based prompting
- Order-based prompting
- V4CharacterCaption: Character-specific prompting with positioning
"""
```
Note: Basic img2img functionality works with V4 preview. For inpainting, the node will automatically use V3 model but can still work on V4-generated images. Vibe transfer will be supported once V4 fully releases.
**Note:** Basic `img2img`, `vibe transfer` and `inpainting` functionality works with V4/V4.5.
nodes.py
+297 -231
View File
@@ -3,11 +3,93 @@ import io
from pathlib import Path
import folder_paths
import zipfile
import json as _json
import copy as _copy
from .utils import *
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
import torch
import numpy as np
from PIL import Image as PILImage
TOOLTIP_LIMIT_OPUS_FREE = "Limit image size and steps for free generation by Opus."
# ------------------------------------------------------------------
# Helper utilities
# ------------------------------------------------------------------
# Accepted canvas sizes (per CR guidance); we will letterbox/pad to one of these
ACCEPTED_CR_SIZES = [(1024, 1536), (1536, 1024), (1472, 1472)]
def _get_user_data(access_token, timeout=120, retry=3):
"""Fetches user data to check Anlas balance. Now a global helper."""
USER_API_BASE_URL = "https://api.novelai.net"
req_mod = requests
if retry is not None and retry > 1:
retries = Retry(
total=retry,
backoff_factor=1,
status_forcelist=[429, 500, 502, 503, 504],
allowed_methods=["GET", "POST"]
)
session = requests.Session()
session.mount("https://", HTTPAdapter(max_retries=retries))
req_mod = session
response = req_mod.get(
f"{USER_API_BASE_URL}/user/data",
headers={"Authorization": f"Bearer {access_token}"},
timeout=timeout
)
response.raise_for_status()
return response.json()
def _choose_cr_canvas(w, h):
"""Select the accepted CR canvas size whose aspect ratio is closest to the source image."""
aspect = w / h
best = None
best_diff = 9e9
for cw, ch in ACCEPTED_CR_SIZES:
diff = abs((cw / ch) - aspect)
if diff < best_diff:
best_diff = diff
best = (cw, ch)
return best
def pad_image_to_canvas(tensor_image, target_size):
"""
Letterbox the given tensor image [1,H,W,C] into target_size (W,H) with black padding,
preserving aspect ratio.
"""
_, H, W, C = tensor_image.shape
tw, th = target_size
arr = (tensor_image[0].cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
mode = "RGBA" if (C == 4) else "RGB"
pil = PILImage.fromarray(arr)
scale = min(tw / W, th / H)
new_w = max(1, int(W * scale))
new_h = max(1, int(H * scale))
pil_resized = pil.resize((new_w, new_h), PILImage.LANCZOS)
if mode == "RGBA":
canvas = PILImage.new("RGBA", (tw, th), (0, 0, 0, 0))
else:
canvas = PILImage.new("RGB", (tw, th), (0, 0, 0))
offset = ((tw - new_w) // 2, (th - new_h) // 2)
canvas.paste(pil_resized, offset)
out = np.array(canvas).astype(np.float32) / 255.0
return torch.from_numpy(out)[None,]
# -------------------------------------------------
# Core simple prompt conversion / utility nodes
# -------------------------------------------------
class PromptToNAID:
@classmethod
def INPUT_TYPES(s):
@@ -41,7 +123,15 @@ class ModelOption:
def INPUT_TYPES(s):
return {
"required": {
"model": (["nai-diffusion-2", "nai-diffusion-furry-3", "nai-diffusion-3", "nai-diffusion-4-curated-preview", "nai-diffusion-4-full", "nai-diffusion-4-5-curated", "nai-diffusion-4-5-full"], { "default": "nai-diffusion-4-5-full" }),
"model": ([
"nai-diffusion-2",
"nai-diffusion-furry-3",
"nai-diffusion-3",
"nai-diffusion-4-curated-preview",
"nai-diffusion-4-full",
"nai-diffusion-4-5-curated",
"nai-diffusion-4-5-full"
], { "default": "nai-diffusion-4-5-full" }),
},
"optional": { "option": ("NAID_OPTION",) },
}
@@ -110,7 +200,6 @@ class VibeTransferOption:
option = copy.deepcopy(option) if option else {}
if "vibe" not in option:
option["vibe"] = []
option["vibe"].append((image, information_extracted, strength))
return (option,)
@@ -135,6 +224,39 @@ class NetworkOption:
option["retry"] = retry
return (option,)
# -------------------------------------------------
# Character Reference (Single Image)
# -------------------------------------------------
class CharacterReferenceOption:
INFO_EXTRACT_DEFAULT = 1.0
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"style_aware": ("BOOLEAN", {"default": True, "tooltip": "Copy style along with identity."}),
"fidelity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "display": "number", "tooltip": "How strictly to match the character (and style if enabled)."}),
},
"optional": {"option": ("NAID_OPTION",),}
}
RETURN_TYPES = ("NAID_OPTION",)
FUNCTION = "set_option"
CATEGORY = "NovelAI"
def set_option(self, image, style_aware, fidelity, option=None):
option = copy.deepcopy(option) if option else {}
fidelity = max(0.0, min(1.0, fidelity))
option["character_reference_single"] = {
"image": image,
"style_aware": style_aware,
"fidelity": fidelity,
"info_extracted": self.INFO_EXTRACT_DEFAULT,
}
return (option,)
# -------------------------------------------------
# Generation Node
# -------------------------------------------------
class GenerateNAID:
def __init__(self):
@@ -169,55 +291,55 @@ class GenerateNAID:
FUNCTION = "generate"
CATEGORY = "NovelAI"
def generate(self, limit_opus_free, width, height, positive, negative, steps, cfg, decrisper, variety, smea, sampler, scheduler, seed, uncond_scale, cfg_rescale, keep_alpha, option=None):
width, height = calculate_resolution(width*height, (width, height))
@staticmethod
def _post_image(access_token, prompt, model, action, parameters, timeout=None, retry=None):
data = {"input": prompt, "model": model, "action": action, "parameters": parameters}
req_mod = requests
if retry is not None and retry > 1:
retries = Retry(total=retry, backoff_factor=1, status_forcelist=[429, 500, 502, 503, 504], allowed_methods=["POST"])
session = requests.Session()
session.mount("https://", HTTPAdapter(max_retries=retries))
req_mod = session
response = req_mod.post(f"{BASE_URL}/ai/generate-image", json=data, headers={"Authorization": f"Bearer {access_token}"}, timeout=timeout)
if response.status_code >= 400:
print("RAW ERROR STATUS:", response.status_code)
print("RAW ERROR BODY:", response.text)
try:
dbg = _copy.deepcopy(data)
p = dbg.get("parameters", {})
if "director_reference_images" in p: p["director_reference_images"] = [i[:60] + "...(trunc)" for i in p["director_reference_images"]]
if "reference_image_multiple" in p: p["reference_image_multiple"] = [i[:60] + "...(trunc)" for i in p["reference_image_multiple"]]
dbg["parameters"] = p
print("OUTGOING PAYLOAD (sanitized):", _json.dumps(dbg)[:2000])
except Exception as e:
print("Payload debug failed:", e)
response.raise_for_status()
return response.content
def generate(self, limit_opus_free, width, height, positive, negative,
steps, cfg, decrisper, variety, smea, sampler, scheduler,
seed, uncond_scale, cfg_rescale, keep_alpha, 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,
"qualityToggle": False,
"params_version": 1, "width": width, "height": height, "scale": cfg, "sampler": sampler, "steps": steps,
"seed": seed, "n_samples": 1, "ucPreset": 3, "qualityToggle": False,
"sm": (smea == "SMEA" or smea == "SMEA+DYN") and sampler != "ddim",
"sm_dyn": smea == "SMEA+DYN" and sampler != "ddim",
"dynamic_thresholding": decrisper,
"skip_cfg_above_sigma": None,
"controlnet_strength": 1.0,
"legacy": False,
"add_original_image": False,
"cfg_rescale": cfg_rescale,
"noise_schedule": scheduler,
"legacy_v3_extend": False,
"uncond_scale": uncond_scale,
"negative_prompt": negative,
"prompt": positive,
"reference_image_multiple": [],
"reference_information_extracted_multiple": [],
"reference_strength_multiple": [],
"sm_dyn": (smea == "SMEA+DYN") and sampler != "ddim",
"dynamic_thresholding": decrisper, "controlnet_strength": 1.0, "legacy": False, "add_original_image": False,
"cfg_rescale": cfg_rescale, "noise_schedule": scheduler, "legacy_v3_extend": False,
"uncond_scale": uncond_scale, "negative_prompt": negative, "prompt": positive,
"reference_image_multiple": [], "reference_information_extracted_multiple": [], "reference_strength_multiple": [],
"extra_noise_seed": seed,
"v4_prompt": {
"use_coords": False,
"use_order": False,
"caption": {
"base_caption": positive,
"char_captions": []
}
},
"v4_negative_prompt": {
"use_coords": False,
"use_order": False,
"caption": {
"base_caption": negative,
"char_captions": []
}
}
"v4_prompt": {"use_coords": False, "use_order": False, "caption": {"base_caption": positive, "char_captions": []}},
"v4_negative_prompt": {"use_coords": False, "use_order": False, "caption": {"base_caption": negative, "char_captions": []}}
}
model = "nai-diffusion-4-5-full"
action = "generate"
@@ -241,106 +363,117 @@ class GenerateNAID:
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))))
vimg, information_extracted, strength = vibe
params["reference_image_multiple"].append(image_to_base64(resize_image(vimg, (width, height))))
params["reference_information_extracted_multiple"].append(information_extracted)
params["reference_strength_multiple"].append(strength)
if "model" in option:
model = option["model"]
if "model" in option: model = option["model"]
if "v4_prompt" in option: params["v4_prompt"].update(option["v4_prompt"])
# Handle V4 options
if "v4_prompt" in option:
params["v4_prompt"].update(option["v4_prompt"])
if "character_reference_single" in option:
ref = option["character_reference_single"]
base_caption = "character&style" if ref["style_aware"] else "character"
ref_img = ref["image"]
_, h_raw, w_raw, _ = ref_img.shape
canvas_w, canvas_h = _choose_cr_canvas(w_raw, h_raw)
padded = pad_image_to_canvas(ref_img, (canvas_w, canvas_h))
params["director_reference_images"] = [image_to_base64(padded)]
params["director_reference_descriptions"] = [{"use_coords": False, "use_order": False, "legacy_uc": False, "caption": {"base_caption": base_caption, "char_captions": []}}]
params["director_reference_strength_values"] = [1.0]
params["director_reference_secondary_strength_values"] = [1.0 - ref["fidelity"]]
params["director_reference_information_extracted"] = [1.0]
timeout = option["timeout"] if option and "timeout" in option else None
retry = option["retry"] if option and "retry" in option else None
timeout = option.get("timeout", 120) if option else 120
retry = option.get("retry", 3) if option else 3
if limit_opus_free:
pixel_limit = 1024*1024
pixel_limit = 1024 * 1024
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
params["width"], params["height"] = calculate_resolution(pixel_limit, (width, height))
if steps > 28: params["steps"] = 28
if variety:
params["skip_cfg_above_sigma"] = calculate_skip_cfg_above_sigma(params["width"], params["height"])
if sampler == "ddim" and model not in ("nai-diffusion-2"):
params["sampler"] = "ddim_v3"
if action == "infill" and model not in ("nai-diffusion-2"):
model = f"{model}-inpainting"
if variety: params["skip_cfg_above_sigma"] = calculate_skip_cfg_above_sigma(params["width"], params["height"])
if sampler == "ddim" and "nai-diffusion-2" not in model: params["sampler"] = "ddim_v3"
if action == "infill" and "nai-diffusion-2" not in model: model = f"{model}-inpainting"
start_anlas = None
try:
user_data = _get_user_data(self.access_token, timeout, retry)
start_anlas = user_data.get("subscription", {}).get("trainingStepsLeft")
if start_anlas is not None: print(f"[NovelAI] Anlas (pre-gen): {start_anlas}")
except Exception as e: print(f"[NovelAI] Anlas tracking failed (pre-gen): {e}")
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
zipped_bytes = self._post_image(self.access_token, positive, model, action, params, timeout, retry)
with zipfile.ZipFile(io.BytesIO(zipped_bytes)) as zipped:
image_bytes = zipped.read(zipped.infolist()[0])
## 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)
full_output_folder, filename, counter, _, _ = 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)
if start_anlas is not None:
try:
user_data_final = _get_user_data(self.access_token, timeout, retry)
final_anlas = user_data_final.get("subscription", {}).get("trainingStepsLeft")
if final_anlas is not None:
print(f"[NovelAI] Generation cost: {start_anlas - final_anlas} Anlas")
print(f"[NovelAI] Anlas (post-gen): {final_anlas}")
except Exception as e: print(f"[NovelAI] Anlas tracking failed (post-gen): {e}")
image = bytes_to_image(image_bytes, keep_alpha)
except Exception as e:
if "ignore_errors" in option and option["ignore_errors"]:
print("ignore error:", e)
else:
raise e
if option and option.get("ignore_errors", False): print("ignore error:", e)
else: raise e
return (image,)
# -------------------------------------------------
# Director Tool Augment Nodes
# -------------------------------------------------
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 limit_opus_free:
pixel_limit = 1024 * 1024
if w * h > pixel_limit:
w, h = calculate_resolution(pixel_limit, (w, h))
w, h = image.shape[2], image.shape[1]
if limit_opus_free and w * h > 1024 * 1024:
w, h = calculate_resolution(1024 * 1024, (w, h))
start_anlas = None
try:
user_data = _get_user_data(access_token)
start_anlas = user_data.get("subscription", {}).get("trainingStepsLeft")
if start_anlas is not None: print(f"[NovelAI] Anlas (pre-augment): {start_anlas}")
except Exception as e: print(f"[NovelAI] Anlas tracking failed (pre-augment): {e}")
base64_image = image_to_base64(resize_image(image, (w, h)))
result_image = blank_image()
try:
# Build request based on NAI v4 API spec
request = {
"image": base64_image,
"req_type": req_type,
"width": w,
"height": h
}
# Add optional parameters if provided
if options:
if "defry" in options:
request["defry"] = options["defry"]
if "prompt" in options:
request["prompt"] = options["prompt"]
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
with zipfile.ZipFile(io.BytesIO(zipped_bytes)) as zipped:
image_bytes = zipped.read(zipped.infolist()[0])
## 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)
full_output_folder, filename, counter, _, _ = 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)
if start_anlas is not None:
try:
user_data_final = _get_user_data(access_token)
final_anlas = user_data_final.get("subscription", {}).get("trainingStepsLeft")
if final_anlas is not None:
print(f"[NovelAI] Augment cost: {start_anlas - final_anlas} Anlas")
print(f"[NovelAI] Anlas (post-augment): {final_anlas}")
except Exception as e: print(f"[NovelAI] Anlas tracking failed (post-augment): {e}")
result_image = bytes_to_image(image_bytes)
except Exception as e:
if ignore_errors:
print("ignore error:", e)
else:
raise e
if ignore_errors: print("ignore error:", e)
else: raise e
return (result_image,)
@@ -350,13 +483,7 @@ class RemoveBGAugment:
self.output_dir = folder_paths.get_output_directory()
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"limit_opus_free": ("BOOLEAN", { "default": True, "tooltip": TOOLTIP_LIMIT_OPUS_FREE }),
"ignore_errors": ("BOOLEAN", { "default": False }),
},
}
return {"required": {"image": ("IMAGE",), "limit_opus_free": ("BOOLEAN", { "default": True, "tooltip": TOOLTIP_LIMIT_OPUS_FREE }), "ignore_errors": ("BOOLEAN", { "default": False }),}}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "augment"
CATEGORY = "NovelAI/director_tools"
@@ -369,13 +496,7 @@ class LineArtAugment:
self.output_dir = folder_paths.get_output_directory()
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"limit_opus_free": ("BOOLEAN", { "default": True, "tooltip": TOOLTIP_LIMIT_OPUS_FREE }),
"ignore_errors": ("BOOLEAN", { "default": False }),
},
}
return {"required": {"image": ("IMAGE",), "limit_opus_free": ("BOOLEAN", { "default": True, "tooltip": TOOLTIP_LIMIT_OPUS_FREE }), "ignore_errors": ("BOOLEAN", { "default": False }),}}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "augment"
CATEGORY = "NovelAI/director_tools"
@@ -388,13 +509,7 @@ class SketchAugment:
self.output_dir = folder_paths.get_output_directory()
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"limit_opus_free": ("BOOLEAN", { "default": True, "tooltip": TOOLTIP_LIMIT_OPUS_FREE }),
"ignore_errors": ("BOOLEAN", { "default": False }),
},
}
return {"required": {"image": ("IMAGE",), "limit_opus_free": ("BOOLEAN", { "default": True, "tooltip": TOOLTIP_LIMIT_OPUS_FREE }), "ignore_errors": ("BOOLEAN", { "default": False }),}}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "augment"
CATEGORY = "NovelAI/director_tools"
@@ -407,15 +522,7 @@ class ColorizeAugment:
self.output_dir = folder_paths.get_output_directory()
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"limit_opus_free": ("BOOLEAN", { "default": True, "tooltip": TOOLTIP_LIMIT_OPUS_FREE }),
"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 {"required": {"image": ("IMAGE",), "limit_opus_free": ("BOOLEAN", { "default": True, "tooltip": TOOLTIP_LIMIT_OPUS_FREE }), "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"
@@ -426,24 +533,10 @@ 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, "tooltip": TOOLTIP_LIMIT_OPUS_FREE }),
"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 {"required": {"image": ("IMAGE",), "limit_opus_free": ("BOOLEAN", { "default": True, "tooltip": TOOLTIP_LIMIT_OPUS_FREE }), "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"
@@ -458,103 +551,72 @@ class DeclutterAugment:
self.output_dir = folder_paths.get_output_directory()
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"limit_opus_free": ("BOOLEAN", { "default": True, "tooltip": TOOLTIP_LIMIT_OPUS_FREE }),
"ignore_errors": ("BOOLEAN", { "default": False }),
},
}
return {"required": {"image": ("IMAGE",), "limit_opus_free": ("BOOLEAN", { "default": True, "tooltip": TOOLTIP_LIMIT_OPUS_FREE }), "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)
# -------------------------------------------------
# Anlas Tracker (Visual Node)
# -------------------------------------------------
class AnlasTrackerNAID:
def __init__(self):
self.access_token = get_access_token()
@classmethod
def INPUT_TYPES(s):
return {
"required": {},
"optional": { "trigger": ("*",) } # Allows chaining to control execution order
}
RETURN_TYPES = ("INT", "STRING",)
RETURN_NAMES = ("anlas_int", "anlas_string",)
FUNCTION = "get_anlas"
CATEGORY = "NovelAI/utils"
def get_anlas(self, trigger=None):
anlas_count = 0
try:
user_data = _get_user_data(self.access_token)
anlas_count = user_data.get("subscription", {}).get("trainingStepsLeft", 0)
print(f"[NovelAI] Current Anlas Balance: {anlas_count}")
except Exception as e:
print(f"[NovelAI] Failed to fetch Anlas balance: {e}")
return (0, "Error fetching Anlas")
return (anlas_count, f"{anlas_count} Anlas")
# -------------------------------------------------
# V4 Base / Negative Prompt nodes
# -------------------------------------------------
class V4BasePrompt:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"base_caption": ("STRING", { "multiline": True }),
}
}
RETURN_TYPES = ("STRING",) # Changed from NAID_OPTION to STRING
FUNCTION = "convert" # Changed from set_option to convert
return {"required": {"base_caption": ("STRING", { "multiline": True }),}}
RETURN_TYPES = ("STRING",)
FUNCTION = "convert"
CATEGORY = "NovelAI/v4"
def convert(self, base_caption):
return (base_caption,) # Simply returns the caption as a string
return (base_caption,)
"""class V4PromptConfig:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"use_coords": ("BOOLEAN", { "default": False }),
"use_order": ("BOOLEAN", { "default": False }),
},
"optional": { "option": ("NAID_OPTION",) },
}
RETURN_TYPES = ("NAID_OPTION",)
FUNCTION = "set_option"
CATEGORY = "NovelAI/v4"
def set_option(self, use_coords, use_order, option=None):
option = copy.deepcopy(option) if option else {}
if "v4_prompt" not in option:
option["v4_prompt"] = {}
option["v4_prompt"]["use_coords"] = use_coords
option["v4_prompt"]["use_order"] = use_order
return (option,)
class V4CharacterCaption:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"char_caption": ("STRING", { "multiline": True }),
"x": ("FLOAT", { "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01 }),
"y": ("FLOAT", { "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01 }),
},
"optional": { "option": ("NAID_OPTION",) },
}
RETURN_TYPES = ("NAID_OPTION",)
FUNCTION = "set_option"
CATEGORY = "NovelAI/v4"
def set_option(self, char_caption, x, y, option=None):
option = copy.deepcopy(option) if option else {}
if "v4_prompt" not in option:
option["v4_prompt"] = {
"caption": {
"base_caption": "",
"char_captions": []
}
}
char_caption_obj = {
"char_caption": char_caption,
"centers": [{"x": x, "y": y}]
}
option["v4_prompt"]["caption"]["char_captions"].append(char_caption_obj)
return (option,)
"""
class V4NegativePrompt:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"negative_caption": ("STRING", { "multiline": True }),
}
}
return {"required": {"negative_caption": ("STRING", { "multiline": True }),}}
RETURN_TYPES = ("STRING",)
FUNCTION = "convert"
CATEGORY = "NovelAI/v4"
def convert(self, negative_caption):
return (negative_caption,)
# -------------------------------------------------
# Registration
# -------------------------------------------------
NODE_CLASS_MAPPINGS = {
"GenerateNAID": GenerateNAID,
@@ -563,6 +625,8 @@ NODE_CLASS_MAPPINGS = {
"InpaintingOptionNAID": InpaintingOption,
"VibeTransferOptionNAID": VibeTransferOption,
"NetworkOptionNAID": NetworkOption,
"CharacterReferenceOptionNAID": CharacterReferenceOption,
"AnlasTrackerNAID": AnlasTrackerNAID, # New node
"MaskImageToNAID": ImageToNAIMask,
"PromptToNAID": PromptToNAID,
"RemoveBGNAID": RemoveBGAugment,
@@ -582,6 +646,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"InpaintingOptionNAID": "InpaintingOption ✒️🅝🅐🅘",
"VibeTransferOptionNAID": "VibeTransferOption ✒️🅝🅐🅘",
"NetworkOptionNAID": "NetworkOption ✒️🅝🅐🅘",
"CharacterReferenceOptionNAID": "Character Reference ✒️🅝🅐🅘",
"AnlasTrackerNAID": "Anlas Tracker ✒️🅝🅐🅘", # New node
"MaskImageToNAID": "Convert Mask Image ✒️🅝🅐🅘",
"PromptToNAID": "Convert Prompt ✒️🅝🅐🅘",
"RemoveBGNAID": "Remove BG ✒️🅝🅐🅘",