Merge pull request #25 from agentd00nut/bugfix/key-valuerror-after-copy

Fix key/value error after copy for ImagePreset
Aedial · Dec 3, 2023 · d652e99 · d652e99
2 parents 37f2250 + 52df5f2
commit d652e99
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 6 deletions.
diff --git a/novelai_api/ImagePreset.py b/novelai_api/ImagePreset.py
@@ -199,6 +199,7 @@ class ImagePreset:
     }
 
     _TYPE_MAPPING = {
+        "legacy": bool,
         "quality_toggle": bool,
         "resolution": (ImageResolution, tuple),
         "uc_preset": (UCPreset, NoneType),
@@ -208,6 +209,7 @@ class ImagePreset:
         "noise": (int, float),
         "strength": (int, float),
         "scale": (int, float),
+        "uncond_scale": (int, float),
         "steps": int,
         "uc": str,
         "smea": bool,
@@ -219,6 +221,8 @@ class ImagePreset:
         "decrisper": bool,
         "add_original_image": bool,
         "mask": str,
+        "cfg_rescale": float,
+        "noise_schedule": str,
     }
 
     # type completion for __setitem__ and __getitem__

diff --git a/novelai_api/_high_level.py b/novelai_api/_high_level.py
@@ -298,12 +298,17 @@ async def _generate(
         global_settings.rep_pen_whitelist = repetition_penalty_default_whitelist
 
         params = {
-            "repetition_penalty_whitelist": list(set(
-                item for sublist in [
-                    global_params.pop("repetition_penalty_whitelist", []),
-                    preset_params.pop("repetition_penalty_whitelist", []),
-                ] for inner_list in sublist for item in inner_list
-            ))
+            "repetition_penalty_whitelist": list(
+                set(
+                    item
+                    for sublist in [
+                        global_params.pop("repetition_penalty_whitelist", []),
+                        preset_params.pop("repetition_penalty_whitelist", []),
+                    ]
+                    for inner_list in sublist
+                    for item in inner_list
+                )
+            )
         }
 
         params.update(preset_params)

diff --git a/tests/api/test_imagegen_samplers.py b/tests/api/test_imagegen_samplers.py
@@ -41,6 +41,7 @@ async def test_samplers(
     logger.info(f"Testing model {model} with sampler {sampler}")
 
     preset = ImagePreset(sampler=sampler)
+    preset.copy()
 
     # Furry doesn't have UCPreset.Preset_Low_Quality_Bad_Anatomy
     if model is ImageModel.Furry: