From ad6caa1a42bc83c1b0df1c1c1951479135990120 Mon Sep 17 00:00:00 2001
From: Mooshieblob <blob@mooshieblob.com>
Date: Thu, 9 Oct 2025 00:39:39 +0800
Subject: [PATCH 1/7] Implement character reference image processing

Added character reference handling and image padding utilities.
---
 nodes.py | 338 ++++++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 259 insertions(+), 79 deletions(-)

diff --git a/nodes.py b/nodes.py
index 77b946c..9f5c97f 100644
--- a/nodes.py
+++ b/nodes.py
@@ -3,11 +3,68 @@ 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 (local) for Character Reference image preparation
+# ------------------------------------------------------------------
+
+# Accepted canvas sizes (per CR guidance); we will letterbox/pad to one of these
+ACCEPTED_CR_SIZES = [(1024, 1536), (1536, 1024), (1472, 1472)]
+
+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 +98,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 +175,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 +199,73 @@ class NetworkOption:
         option["retry"] = retry
         return (option,)
 
+# -------------------------------------------------
+# Character Reference (Single Image)
+# -------------------------------------------------
+
+class CharacterReferenceOption:
+    """
+    Single-image Character Reference node.
+
+    API requirement (observed): director_reference_information_extracted must be EXACTLY 1.0.
+    Inputs:
+      - image
+      - style_aware
+      - fidelity (0-1)
+
+    Mapping:
+      style_aware False:
+        base_caption = "character"
+        primary_strength = fidelity
+        secondary_strength = 0.0
+      style_aware True:
+        base_caption = "character&style"
+        primary_strength = fidelity
+        secondary_strength = fidelity
+    """
+    INFO_EXTRACT_DEFAULT = 1.0  # Required by backend
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "style_aware": ("BOOLEAN", {
+                    "default": True,
+                    "tooltip": "Copy style along with identity."
+                }),
+                "fidelity": ("FLOAT", {
+                    "default": 0.65,
+                    "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,10 +300,53 @@ 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,
@@ -185,9 +359,9 @@ class GenerateNAID:
             "ucPreset": 3,
             "qualityToggle": False,
             "sm": (smea == "SMEA" or smea == "SMEA+DYN") and sampler != "ddim",
-            "sm_dyn": smea == "SMEA+DYN" and sampler != "ddim",
+            "sm_dyn": (smea == "SMEA+DYN") and sampler != "ddim",
             "dynamic_thresholding": decrisper,
-            "skip_cfg_above_sigma": None,
+            # skip_cfg_above_sigma added later only if variety True
             "controlnet_strength": 1.0,
             "legacy": False,
             "add_original_image": False,
@@ -218,6 +392,7 @@ class GenerateNAID:
                 }
             }
         }
+
         model = "nai-diffusion-4-5-full"
         action = "generate"
 
@@ -241,23 +416,69 @@ 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"]
 
-            # Handle V4 options
             if "v4_prompt" in option:
                 params["v4_prompt"].update(option["v4_prompt"])
 
+            # --- CHARACTER REFERENCE SINGLE INJECTION ---
+            if "character_reference_single" in option:
+                ref = option["character_reference_single"]
+                style_aware = ref["style_aware"]
+                fidelity = ref["fidelity"]
+                # Backend requires EXACT 1.0
+                info_extracted = 1.0
+
+                # Determine base_caption & strengths
+                if style_aware:
+                    base_caption = "character&style"
+                    primary_strength = fidelity
+                    secondary_strength = fidelity
+                else:
+                    base_caption = "character"
+                    primary_strength = fidelity
+                    secondary_strength = fidelity
+
+                # Prepare padded image to accepted CR canvas
+                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))
+                b64_img = image_to_base64(padded)
+
+                params["director_reference_images"] = [b64_img]
+                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"] = [primary_strength]
+                params["director_reference_secondary_strength_values"] = [secondary_strength]
+                params["director_reference_information_extracted"] = [info_extracted]
+
+                print("[CR DEBUG] Injected CR:",
+                      "base_caption=", base_caption,
+                      "primary=", primary_strength,
+                      "secondary=", secondary_strength,
+                      "info_extracted=", info_extracted,
+                      "canvas=", (canvas_w, canvas_h))
+            # --- END CHARACTER REFERENCE SINGLE INJECTION ---
+
         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
+            pixel_limit = 1024 * 1024
             if width * height > pixel_limit:
                 max_width, max_height = calculate_resolution(pixel_limit, (width, height))
                 params["width"] = max_width
@@ -276,12 +497,14 @@ class GenerateNAID:
 
         image = blank_image()
         try:
-            zipped_bytes = generate_image(self.access_token, positive, model, action, params, timeout, retry)
+            zipped_bytes = self._post_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
+            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)
+            # 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)
@@ -289,13 +512,16 @@ class GenerateNAID:
 
             image = bytes_to_image(image_bytes, keep_alpha)
         except Exception as e:
-            if "ignore_errors" in option and option["ignore_errors"]:
+            if option and "ignore_errors" in option and option["ignore_errors"]:
                 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)
@@ -309,7 +535,6 @@ def base_augment(access_token, output_dir, limit_opus_free, ignore_errors, req_t
     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,
@@ -317,7 +542,6 @@ def base_augment(access_token, output_dir, limit_opus_free, ignore_errors, req_t
             "height": h
         }
 
-        # Add optional parameters if provided
         if options:
             if "defry" in options:
                 request["defry"] = options["defry"]
@@ -326,10 +550,11 @@ def base_augment(access_token, output_dir, limit_opus_free, ignore_errors, req_t
 
         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
+        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, 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)
@@ -470,6 +695,11 @@ class DeclutterAugment:
     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)
+
+# -------------------------------------------------
+# V4 Base / Negative Prompt nodes
+# -------------------------------------------------
+
 class V4BasePrompt:
     @classmethod
     def INPUT_TYPES(s):
@@ -479,67 +709,12 @@ class V4BasePrompt:
             }
         }
 
-    RETURN_TYPES = ("STRING",)  # Changed from NAID_OPTION to STRING
-    FUNCTION = "convert"        # Changed from set_option to convert
+    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):
@@ -555,6 +730,9 @@ class V4NegativePrompt:
     def convert(self, negative_caption):
         return (negative_caption,)
 
+# -------------------------------------------------
+# Registration
+# -------------------------------------------------
 
 NODE_CLASS_MAPPINGS = {
     "GenerateNAID": GenerateNAID,
@@ -563,6 +741,7 @@ NODE_CLASS_MAPPINGS = {
     "InpaintingOptionNAID": InpaintingOption,
     "VibeTransferOptionNAID": VibeTransferOption,
     "NetworkOptionNAID": NetworkOption,
+    "CharacterReferenceOptionNAID": CharacterReferenceOption,
     "MaskImageToNAID": ImageToNAIMask,
     "PromptToNAID": PromptToNAID,
     "RemoveBGNAID": RemoveBGAugment,
@@ -582,6 +761,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "InpaintingOptionNAID": "InpaintingOption ✒️🅝🅐🅘",
     "VibeTransferOptionNAID": "VibeTransferOption ✒️🅝🅐🅘",
     "NetworkOptionNAID": "NetworkOption ✒️🅝🅐🅘",
+    "CharacterReferenceOptionNAID": "Character Reference ✒️🅝🅐🅘",
     "MaskImageToNAID": "Convert Mask Image ✒️🅝🅐🅘",
     "PromptToNAID": "Convert Prompt ✒️🅝🅐🅘",
     "RemoveBGNAID": "Remove BG ✒️🅝🅐🅘",

From a0ac51f09c44cdb0b3eb78c1d87883e334edcbc3 Mon Sep 17 00:00:00 2001
From: Mooshieblob <blob@mooshieblob.com>
Date: Thu, 9 Oct 2025 02:11:54 +0800
Subject: [PATCH 2/7] Update fidelity default value and related calculations

---
 nodes.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/nodes.py b/nodes.py
index 9f5c97f..b17ac25 100644
--- a/nodes.py
+++ b/nodes.py
@@ -235,7 +235,7 @@ class CharacterReferenceOption:
                     "tooltip": "Copy style along with identity."
                 }),
                 "fidelity": ("FLOAT", {
-                    "default": 0.65,
+                    "default": 1.0,
                     "min": 0.0,
                     "max": 1.0,
                     "step": 0.01,
@@ -432,18 +432,15 @@ class GenerateNAID:
                 ref = option["character_reference_single"]
                 style_aware = ref["style_aware"]
                 fidelity = ref["fidelity"]
-                # Backend requires EXACT 1.0
                 info_extracted = 1.0
+                primary_strength = 1.0
+                secondary_strength = 1.0 - fidelity
 
                 # Determine base_caption & strengths
                 if style_aware:
                     base_caption = "character&style"
-                    primary_strength = fidelity
-                    secondary_strength = fidelity
                 else:
                     base_caption = "character"
-                    primary_strength = fidelity
-                    secondary_strength = fidelity
 
                 # Prepare padded image to accepted CR canvas
                 ref_img = ref["image"]

From 44d65b286bbc3dfc21e6426c95630b69402a3815 Mon Sep 17 00:00:00 2001
From: Mooshieblob <blob@mooshieblob.com>
Date: Thu, 9 Oct 2025 03:36:56 +0800
Subject: [PATCH 3/7] Add user data fetching and Anlas tracking features

Added a new helper function to fetch user data and integrated Anlas tracking into the augment and generation processes. Refactored some sections for clarity and efficiency.
---
 nodes.py | 401 ++++++++++++++++++++-----------------------------------
 1 file changed, 145 insertions(+), 256 deletions(-)

diff --git a/nodes.py b/nodes.py
index b17ac25..cec25de 100644
--- a/nodes.py
+++ b/nodes.py
@@ -17,12 +17,37 @@ from PIL import Image as PILImage
 TOOLTIP_LIMIT_OPUS_FREE = "Limit image size and steps for free generation by Opus."
 
 # ------------------------------------------------------------------
-# Helper utilities (local) for Character Reference image preparation
+# 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
@@ -204,54 +229,20 @@ class NetworkOption:
 # -------------------------------------------------
 
 class CharacterReferenceOption:
-    """
-    Single-image Character Reference node.
-
-    API requirement (observed): director_reference_information_extracted must be EXACTLY 1.0.
-    Inputs:
-      - image
-      - style_aware
-      - fidelity (0-1)
-
-    Mapping:
-      style_aware False:
-        base_caption = "character"
-        primary_strength = fidelity
-        secondary_strength = 0.0
-      style_aware True:
-        base_caption = "character&style"
-        primary_strength = fidelity
-        secondary_strength = fidelity
-    """
-    INFO_EXTRACT_DEFAULT = 1.0  # Required by backend
-
+    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)."
-                }),
+                "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",),
-            }
+            "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))
@@ -306,22 +297,12 @@ class GenerateNAID:
 
         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"]
-            )
+            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
-        )
+        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)
@@ -329,10 +310,8 @@ class GenerateNAID:
             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"]]
+                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:
@@ -348,49 +327,17 @@ class GenerateNAID:
         width, height = calculate_resolution(width * height, (width, height))
 
         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 added later only if variety True
-            "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": [],
+            "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"
@@ -421,98 +368,67 @@ class GenerateNAID:
                     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"])
 
-            if "v4_prompt" in option:
-                params["v4_prompt"].update(option["v4_prompt"])
-
-            # --- CHARACTER REFERENCE SINGLE INJECTION ---
             if "character_reference_single" in option:
                 ref = option["character_reference_single"]
-                style_aware = ref["style_aware"]
-                fidelity = ref["fidelity"]
-                info_extracted = 1.0
-                primary_strength = 1.0
-                secondary_strength = 1.0 - fidelity
-
-                # Determine base_caption & strengths
-                if style_aware:
-                    base_caption = "character&style"
-                else:
-                    base_caption = "character"
-
-                # Prepare padded image to accepted CR canvas
+                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))
-                b64_img = image_to_base64(padded)
+                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]
 
-                params["director_reference_images"] = [b64_img]
-                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"] = [primary_strength]
-                params["director_reference_secondary_strength_values"] = [secondary_strength]
-                params["director_reference_information_extracted"] = [info_extracted]
-
-                print("[CR DEBUG] Injected CR:",
-                      "base_caption=", base_caption,
-                      "primary=", primary_strength,
-                      "secondary=", secondary_strength,
-                      "info_extracted=", info_extracted,
-                      "canvas=", (canvas_w, canvas_h))
-            # --- END CHARACTER REFERENCE SINGLE INJECTION ---
-
-        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
             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 = self._post_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
+            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 option and "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,)
 
@@ -521,48 +437,43 @@ class GenerateNAID:
 # -------------------------------------------------
 
 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:
-        request = {
-            "image": base64_image,
-            "req_type": req_type,
-            "width": w,
-            "height": h
-        }
-
-        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])
+        with zipfile.ZipFile(io.BytesIO(zipped_bytes)) as zipped:
+            image_bytes = zipped.read(zipped.infolist()[0])
 
-        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,)
 
@@ -572,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"
@@ -591,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"
@@ -610,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"
@@ -629,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"
@@ -648,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"
@@ -680,19 +551,45 @@ 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
 # -------------------------------------------------
@@ -700,12 +597,7 @@ class DeclutterAugment:
 class V4BasePrompt:
     @classmethod
     def INPUT_TYPES(s):
-        return {
-            "required": {
-                "base_caption": ("STRING", { "multiline": True }),
-            }
-        }
-
+        return {"required": {"base_caption": ("STRING", { "multiline": True }),}}
     RETURN_TYPES = ("STRING",)
     FUNCTION = "convert"
     CATEGORY = "NovelAI/v4"
@@ -715,12 +607,7 @@ class V4BasePrompt:
 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"
@@ -739,6 +626,7 @@ NODE_CLASS_MAPPINGS = {
     "VibeTransferOptionNAID": VibeTransferOption,
     "NetworkOptionNAID": NetworkOption,
     "CharacterReferenceOptionNAID": CharacterReferenceOption,
+    "AnlasTrackerNAID": AnlasTrackerNAID, # New node
     "MaskImageToNAID": ImageToNAIMask,
     "PromptToNAID": PromptToNAID,
     "RemoveBGNAID": RemoveBGAugment,
@@ -759,6 +647,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "VibeTransferOptionNAID": "VibeTransferOption ✒️🅝🅐🅘",
     "NetworkOptionNAID": "NetworkOption ✒️🅝🅐🅘",
     "CharacterReferenceOptionNAID": "Character Reference ✒️🅝🅐🅘",
+    "AnlasTrackerNAID": "Anlas Tracker ✒️🅝🅐🅘", # New node
     "MaskImageToNAID": "Convert Mask Image ✒️🅝🅐🅘",
     "PromptToNAID": "Convert Prompt ✒️🅝🅐🅘",
     "RemoveBGNAID": "Remove BG ✒️🅝🅐🅘",

From 58d44dfcfc33c713d0ab1c87ad5d51b838d093e6 Mon Sep 17 00:00:00 2001
From: Mooshieblob <blob@mooshieblob.com>
Date: Thu, 9 Oct 2025 03:45:10 +0800
Subject: [PATCH 4/7] Improve clarity and consistency in README.md

Updated README.md for clarity and consistency in instructions and notes.
---
 README.md | 167 +++++++++++++++++++++---------------------------------
 1 file changed, 66 insertions(+), 101 deletions(-)

diff --git a/README.md b/README.md
index e9efbad..d208148 100644
--- a/README.md
+++ b/README.md
@@ -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) <!-- Placeholder for actual image -->
+
+- **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.
+
+**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` and `inpainting` functionality works with V4/V4.5. `Vibe Transfer` is not officially supported for V4+ models.

From f8fcf354399e25089a8f188ecea7e9484ca3c828 Mon Sep 17 00:00:00 2001
From: Mooshieblob <blob@mooshieblob.com>
Date: Thu, 9 Oct 2025 03:54:42 +0800
Subject: [PATCH 5/7] Fix image syntax and enhance fidelity explanation

Updated image syntax and clarified fidelity description.
---
 README.md | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index d208148..097f999 100644
--- a/README.md
+++ b/README.md
@@ -70,10 +70,10 @@ Connect multiple `VibeTransferOptionNAID` nodes to combine their influences.
 
 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) <!-- Placeholder for actual image -->
+![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.
+-   **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.
 

From 282b176dcb368166e3c9785f817a9a9853c0a0bb Mon Sep 17 00:00:00 2001
From: Mooshieblob <blob@mooshieblob.com>
Date: Thu, 9 Oct 2025 04:08:18 +0800
Subject: [PATCH 6/7] Update README to clarify V4+ model support

Removed mention of unsupported 'Vibe Transfer' for V4+ models.
---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 097f999..d15fa36 100644
--- a/README.md
+++ b/README.md
@@ -151,4 +151,4 @@ V4BasePrompt -----> positive
 V4NegativePrompt -> negative
 ```
 
-**Note:** Basic `img2img` and `inpainting` functionality works with V4/V4.5. `Vibe Transfer` is not officially supported for V4+ models.
+**Note:** Basic `img2img` and `inpainting` functionality works with V4/V4.5.

From 52cdcfb53f1b301a7ddd5158ef0e299284eb8d79 Mon Sep 17 00:00:00 2001
From: Mooshieblob <blob@mooshieblob.com>
Date: Thu, 9 Oct 2025 04:08:47 +0800
Subject: [PATCH 7/7] Update README to include vibe transfer functionality

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index d15fa36..bc45daf 100644
--- a/README.md
+++ b/README.md
@@ -151,4 +151,4 @@ V4BasePrompt -----> positive
 V4NegativePrompt -> negative
 ```
 
-**Note:** Basic `img2img` and `inpainting` functionality works with V4/V4.5.
+**Note:** Basic `img2img`, `vibe transfer` and `inpainting` functionality works with V4/V4.5.