From b98fa38ad20dcc4723dee6f6fb7d087773ba5f3f Mon Sep 17 00:00:00 2001
From: bedovyy <bedovyy@gmail.com>
Date: Fri, 20 Sep 2024 16:59:34 +0900
Subject: [PATCH] add `skip_cfg_above_sigma` on generator

---
 nodes.py       | 7 ++++++-
 pyproject.toml | 2 +-
 utils.py       | 4 ++++
 3 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/nodes.py b/nodes.py
index b971d8a..b7202f5 100644
--- a/nodes.py
+++ b/nodes.py
@@ -151,6 +151,7 @@ class GenerateNAID:
                 "negative": ("STRING", { "default": "lowres", "multiline": True, "dynamicPrompts": False }),
                 "steps": ("INT", { "default": 28, "min": 0, "max": 50, "step": 1, "display": "number" }),
                 "cfg": ("FLOAT", { "default": 5.0, "min": 0.0, "max": 10.0, "step": 0.1, "display": "number" }),
+                "variety" : ("BOOLEAN", { "default": False }),
                 "decrisper": ("BOOLEAN", { "default": False }),
                 "smea": (["none", "SMEA", "SMEA+DYN"], { "default": "none" }),
                 "sampler": (["k_euler", "k_euler_ancestral", "k_dpmpp_2s_ancestral", "k_dpmpp_2m_sde", "k_dpmpp_2m", "k_dpmpp_sde", "ddim"], { "default": "k_euler" }),
@@ -166,7 +167,7 @@ class GenerateNAID:
     FUNCTION = "generate"
     CATEGORY = "NovelAI"
 
-    def generate(self, limit_opus_free, width, height, positive, negative, steps, cfg, decrisper, smea, sampler, scheduler, seed, uncond_scale, cfg_rescale, option=None):
+    def generate(self, limit_opus_free, width, height, positive, negative, steps, cfg, decrisper, variety, smea, sampler, scheduler, seed, uncond_scale, cfg_rescale, option=None):
         width, height = calculate_resolution(width*height, (width, height))
 
         # ref. novelai_api.ImagePreset
@@ -184,6 +185,7 @@ class GenerateNAID:
             "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,
@@ -236,6 +238,9 @@ class GenerateNAID:
             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 == "nai-diffusion-3":
             params["sampler"] = "ddim_v3"
 
diff --git a/pyproject.toml b/pyproject.toml
index df17e3e..1e27e89 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "comfyui_naidgenerator"
 description = "This extension helps generate images through NAI."
-version = "1.0.5"
+version = "1.0.6"
 license = { file = "LICENSE" }
 dependencies = ["python-dotenv", "argon2-cffi"]
 
diff --git a/utils.py b/utils.py
index 3446ecf..3482f82 100644
--- a/utils.py
+++ b/utils.py
@@ -138,6 +138,10 @@ def calculate_resolution(pixel_count, aspect_ratio):
     height = int(np.floor(k * h / w) * 64)
     return width, height
 
+def calculate_skip_cfg_above_sigma(w, h):
+    # 832 * 1216
+    return (w * h / 1011712) ** 0.5 * 19
+
 
 def prompt_to_stack(sentence):
     result = []