[API] Add Vibe feature + new ImagePreset constructor

README_TEMPLATE.md

@@ -4,12 +4,14 @@ Python API for the NovelAI REST API
 This module is intended to be used by developers as a helper for using NovelAI's REST API.
 
 | Category         | Badges                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                        |
-|-----------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+|------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
 | Pypi             | [![PyPI](https://img.shields.io/pypi/v/novelai-api)](https://pypi.org/project/novelai-api) [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/novelai-api)](https://pypi.org/project/novelai-api) [![PyPI - License](https://img.shields.io/pypi/l/novelai-api)](https://pypi.org/project/novelai-api/) [![PyPI - Format](https://img.shields.io/pypi/format/novelai-api)](https://pypi.org/project/novelai-api/)                                                                                                                                                                                                                                                                                               |
 | Quality checking | [![Python package](https://github.com/Aedial/novelai-api/actions/workflows/python-package.yml/badge.svg)](https://github.com/Aedial/novelai-api/actions/workflows/python-package.yml) [![Python package](https://github.com/Aedial/novelai-api/actions/workflows/codeql-analysis.yml/badge.svg)](https://github.com/Aedial/novelai-api/actions/workflows/codeql-analysis.yml) [![linting: pylint](https://img.shields.io/badge/linting-pylint-yellowgreen)](https://github.com/PyCQA/pylint) [![security: bandit](https://img.shields.io/badge/security-bandit-yellow.svg)](https://github.com/PyCQA/bandit) [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black) |
 | Stats            | [![GitHub top language](https://img.shields.io/github/languages/top/Aedial/novelai-api)](https://github.com/Aedial/novelai-api/search?l=python) ![Libraries.io dependency status for GitHub repo](https://img.shields.io/librariesio/github/Aedial/novelai-api) ![GitHub repo size](https://img.shields.io/github/repo-size/Aedial/novelai-api) ![Pypi package size](https://byob.yarr.is/Aedial/novelai-api/pypi-size) ![GitHub issues](https://img.shields.io/github/issues-raw/Aedial/novelai-api) ![GitHub pull requests](https://img.shields.io/github/issues-pr-raw/Aedial/novelai-api)                                                                                                                                 |
-| Activity         | ![GitHub last commit](https://img.shields.io/github/last-commit/Aedial/novelai-api) ![GitHub commits since tagged version](https://img.shields.io/github/commits-since/Aedial/novelai-api/v{bumped_version}) ![GitHub commit activity](https://img.shields.io/github/commit-activity/m/Aedial/novelai-api)                                                                                                                                                                                                                                                                                                                                                                                                                     |
+| Activity         | ![GitHub last commit](https://img.shields.io/github/last-commit/Aedial/novelai-api) ![GitHub commits since tagged version](https://img.shields.io/github/commits-since/Aedial/novelai-api/v{bumped_version}) ![GitHub commit activity](https://img.shields.io/github/commit-activity/m/Aedial/novelai-api)                                                                                                                                                                                                                                                                                                                                                                                                                    |
 
+Retired versions: 3.7.2
+Final commit of retired versions can be found with the tag `py<version>` (e.g. `py3.7.2`).
 
 # Usage
 Download via [pip](https://pypi.org/project/novelai-api):

example/generate_image.py

@@ -21,13 +21,14 @@ async def main():
     async with API() as api_handler:
         api = api_handler.api
 
-        preset = ImagePreset()
+        model = ImageModel.Anime_Full
+        preset = ImagePreset.from_default_config(model)
         preset.seed = 42
 
         # multiple images
         # preset.n_samples = 4
         i = 0
-        async for _, img in api.high_level.generate_image("1girl", ImageModel.Anime_Full, preset):
+        async for _, img in api.high_level.generate_image("1girl", model, preset):
             (d / f"image_1_{i}.png").write_bytes(img)
 
             i += 1

example/generate_image_test_samplers.py

@@ -12,7 +12,7 @@
 from pathlib import Path
 
 from example.boilerplate import API
-from novelai_api.ImagePreset import ImageModel, ImagePreset, ImageSampler
+from novelai_api.ImagePreset import ImageModel, ImagePreset, ImageResolution, ImageSampler, UCPreset
 from novelai_api.NovelAIError import NovelAIError
 
 
@@ -23,14 +23,28 @@ async def main():
     async with API() as api_handler:
         api = api_handler.api
 
-        preset = ImagePreset()
-        preset.seed = 42
+        model = ImageModel.Anime_v3
 
-        for sampler in ImageSampler:
+        preset = ImagePreset.from_default_config(model)
+        preset.resolution = ImageResolution.Normal_Portrait_v3
+        preset.seed = 1796796669
+        preset.scale = 5
+        preset.uc_preset = UCPreset.Preset_None
+        preset.uc = "{{{worst quality, low quality, bad fingers}}},"
+
+        preset.quality_toggle = False
+
+        prompt = (
+            "1girl, smile to viewer, sunny day, frilly white dress, lens flare, sunrays, "
+            "{{detailed fingers, bold outline}}, best quality, amazing quality, very aesthetic, absurdres"
+        )
+        samplers = [ImageSampler.ddim]
+
+        for sampler in samplers:
             preset.sampler = sampler
 
             try:
-                async for _, img in api.high_level.generate_image("1girl", ImageModel.Anime_Full, preset):
+                async for _, img in api.high_level.generate_image(prompt, model, preset):
                     (d / f"image_{sampler.value}").write_bytes(img)
 
                 print(f"Generated with {sampler.value}")

example/generate_image_with_controlnet.py

@@ -28,14 +28,16 @@ async def main():
         controlnet = ControlNetModel.Form_Lock
         _, mask = await api.low_level.generate_controlnet_mask(controlnet, image)
 
-        preset = ImagePreset()
+        model = ImageModel.Anime_Full
+
+        preset = ImagePreset.from_default_config(model)
         preset.controlnet_model = controlnet
         preset.controlnet_condition = base64.b64encode(mask).decode()
         preset.controlnet_strength = 1.5
         preset.seed = 42
 
         # NOTE: for some reasons, the images with controlnet are slightly different
-        async for _, img in api.high_level.generate_image("1girl", ImageModel.Anime_Full, preset):
+        async for _, img in api.high_level.generate_image("1girl", model, preset):
             (d / "image_with_controlnet.png").write_bytes(img)
 
 

example/generate_image_with_img2img.py

@@ -24,16 +24,16 @@ async def main():
 
         image = base64.b64encode((d / "image.png").read_bytes()).decode()
 
-        preset = ImagePreset()
+        model = ImageModel.Anime_Full
+
+        preset = ImagePreset.from_default_config(model)
         preset.noise = 0.1
         # note that steps = 28, not 50, which mean strength needs to be adjusted accordingly
         preset.strength = 0.5
         preset.image = image
         preset.seed = 42
 
-        async for _, img in api.high_level.generate_image(
-            "1girl", ImageModel.Anime_Full, preset, ImageGenerationType.IMG2IMG
-        ):
+        async for _, img in api.high_level.generate_image("1girl", model, preset, ImageGenerationType.IMG2IMG):
             (d / "image_with_img2img.png").write_bytes(img)
 
 

example/generate_image_with_inpainting.py

@@ -25,14 +25,14 @@ async def main():
         image = base64.b64encode((d / "image.png").read_bytes()).decode()
         mask = base64.b64encode((d / "inpainting_mask.png").read_bytes()).decode()
 
-        preset = ImagePreset()
+        model = ImageModel.Inpainting_Anime_Full
+
+        preset = ImagePreset.from_default_config(model)
         preset.image = image
         preset.mask = mask
         preset.seed = 42
 
-        async for _, img in api.high_level.generate_image(
-            "1girl", ImageModel.Inpainting_Anime_Full, preset, ImageGenerationType.INPAINTING
-        ):
+        async for _, img in api.high_level.generate_image("1girl", model, preset, ImageGenerationType.INPAINTING):
             (d / "image_with_inpainting.png").write_bytes(img)
 
 

novelai_api/ImagePreset.py

@@ -6,6 +6,7 @@
 import random
 import sys
 import typing
+from pathlib import Path
 from typing import Any, Dict, Optional, Tuple, Union
 
 from novelai_api.ImagePreset_CostTables import DDIM_COSTS, NAI_COSTS, SMEA_COSTS, SMEA_DYN_COSTS
@@ -21,7 +22,7 @@ class ImageModel(enum.Enum):
     Anime_Full = "nai-diffusion"
     Furry = "nai-diffusion-furry"
 
-    Inainting_Anime_Curated = "safe-diffusion-inpainting"
+    Inpainting_Anime_Curated = "safe-diffusion-inpainting"
     Inpainting_Anime_Full = "nai-diffusion-inpainting"
     Inpainting_Furry = "furry-diffusion-inpainting"
 
@@ -188,7 +189,7 @@ class ImagePreset:
     }
 
     # inpainting presets are the same as the normal ones
-    _UC_Presets[ImageModel.Inainting_Anime_Curated] = _UC_Presets[ImageModel.Anime_Curated]
+    _UC_Presets[ImageModel.Inpainting_Anime_Curated] = _UC_Presets[ImageModel.Anime_Curated]
     _UC_Presets[ImageModel.Inpainting_Anime_Full] = _UC_Presets[ImageModel.Anime_Full]
     _UC_Presets[ImageModel.Inpainting_Furry] = _UC_Presets[ImageModel.Furry]
     _UC_Presets[ImageModel.Inpainting_Anime_v3] = _UC_Presets[ImageModel.Anime_v3]
@@ -209,6 +210,8 @@ class ImagePreset:
     # type completion for __setitem__ and __getitem__
     #: https://docs.novelai.net/image/qualitytags.html
     quality_toggle: bool
+    #: Automatically uses SMEA when image is above 1 megapixel
+    auto_smea: bool
     #: Resolution of the image to generate as ImageResolution or a (width, height) tuple
     resolution: Union[ImageResolution, Tuple[int, int]]
     #: Default UC to prepend to the UC
@@ -253,45 +256,117 @@ class ImagePreset:
     cfg_rescale: float
     #: ??? (TODO: use an enum ? - valid values: native, karras, exponential, polyexponential)
     noise_schedule: str
+    #: b64-encoded png image for Vibe Transfer
+    reference_image: str
+    #: https://docs.novelai.net/.image/vibetransfer.html#information-extracted
+    reference_information_extracted: float
+    #: https://docs.novelai.net/.image/vibetransfer.html#reference-strength
+    reference_strength: float
 
-    _DEFAULT = {
-        "legacy": False,
-        "quality_toggle": True,
-        "resolution": ImageResolution.Normal_Portrait,
-        "uc_preset": UCPreset.Preset_Low_Quality_Bad_Anatomy,
-        "n_samples": 1,
-        "seed": 0,
-        # TODO: set ImageSampler.k_dpmpp_2m as default ?
-        "sampler": ImageSampler.k_euler_ancestral,
-        "steps": 28,
-        "scale": 11,
-        "uncond_scale": 1.0,
-        "uc": "",
-        "smea": False,
-        "smea_dyn": False,
-        "decrisper": False,
-        "controlnet_strength": 1.0,
-        "add_original_image": False,
-        "cfg_rescale": 0.0,
-        "noise_schedule": "native",
-    }
+    #: Use the old behavior of prompt separation at the 75 tokens mark (can cut words in half)
+    legacy_v3_extend: bool
 
     _settings: Dict[str, Any]
 
     #: Seed provided when generating an image with seed 0 (default). Seed is also in metadata, but might be a hassle
     last_seed: int
 
+    @classmethod
+    def from_file(cls, path: Union[str, bytes, os.PathLike, int]) -> "ImagePreset":
+        """
+        Write the preset to a file
+
+        :param path: Path to the file to read the preset from
+        """
+
+        with open(path, encoding="utf-8") as f:
+            data = json.loads(f.read())
+
+        return cls(**data)
+
+    def to_file(self, path: Union[str, bytes, os.PathLike, int]):
+        """
+        Load the preset from a file
+
+        :param path: Path to the file to write the preset to
+        """
+
+        with open(path, "w", encoding="utf-8") as f:
+            f.write(json.dumps(self._settings))
+
     @expand_kwargs(_TYPE_MAPPING.keys(), _TYPE_MAPPING.values())
     def __init__(self, **kwargs):
-        object.__setattr__(self, "_settings", self._DEFAULT.copy())
+        """
+        Create an empty ImagePreset. Use the "from_*_config" functions to create a
+        """
+
+        object.__setattr__(self, "_settings", {})
         self.update(kwargs)
 
         object.__setattr__(self, "last_seed", 0)
 
+    @classmethod
+    def from_v1_config(cls):
+        """
+        Create a new ImagePreset with the default settings from the v1 config
+        """
+
+        return cls.from_file(Path(__file__).parent / "image_presets" / "presets_v1" / "default.preset")
+
+    @classmethod
+    def from_v2_config(cls):
+        """
+        Create a new ImagePreset with the default settings from the v2 config
+        """
+
+        return cls.from_file(Path(__file__).parent / "image_presets" / "presets_v2" / "default.preset")
+
+    @classmethod
+    def from_v3_config(cls):
+        """
+        Create a new ImagePreset with the default settings from the v3 config
+        """
+
+        return cls.from_file(Path(__file__).parent / "image_presets" / "presets_v3" / "default.preset")
+
+    @classmethod
+    def from_default_config(cls, model: ImageModel) -> "ImagePreset":
+        """
+        Create a new ImagePreset with the default settings inferring the version from the model
+
+        :param model: Model to use
+        """
+
+        if model in (
+            ImageModel.Anime_Curated,
+            ImageModel.Anime_Full,
+            ImageModel.Furry,
+            ImageModel.Inpainting_Anime_Curated,
+            ImageModel.Inpainting_Anime_Full,
+            ImageModel.Inpainting_Furry,
+        ):
+            return cls.from_v1_config()
+        elif model in (ImageModel.Anime_v2,):
+            return cls.from_v2_config()
+        elif model in (ImageModel.Anime_v3, ImageModel.Inpainting_Anime_v3):
+            return cls.from_v3_config()
+
     def __setitem__(self, key: str, value: Any):
         if key not in self._TYPE_MAPPING:
             raise ValueError(f"'{key}' is not a valid setting")
 
+        # try to cast into enum if possible
+        types = self._TYPE_MAPPING[key]
+        if not isinstance(types, tuple):
+            types = (types,)
+
+        enum_types = [t for t in types if t.__class__ is enum.EnumMeta]
+        if enum_types and isinstance(value, str):
+            for enum_type in enum_types:
+                if value in enum_type.__members__:  # noqa
+                    value = enum_type[value]  # noqa
+                    break
+
         if not isinstance(value, self._TYPE_MAPPING[key]):  # noqa (pycharm PY-36317)
             raise ValueError(f"Expected type '{self._TYPE_MAPPING[key]}' for {key}, but got type '{type(value)}'")
 
@@ -501,29 +576,6 @@ def calculate_cost(
         opus_discount = is_opus and steps <= 28 and (r <= 640 * 640 if version == 1 else r <= 1024 * 1024)
         return per_sample * (n_samples - int(opus_discount))
 
-    @classmethod
-    def from_file(cls, path: Union[str, bytes, os.PathLike, int]) -> "ImagePreset":
-        """
-        Write the preset to a file
-
-        :param path: Path to the file to read the preset from
-        """
-
-        with open(path, encoding="utf-8") as f:
-            data = json.loads(f.read())
-
-        return cls(**data)
-
-    def to_file(self, path: Union[str, bytes, os.PathLike, int]):
-        """
-        Load the preset from a file
-
-        :param path: Path to the file to write the preset to
-        """
-
-        with open(path, "w", encoding="utf-8") as f:
-            f.write(json.dumps(self._settings))
-
 
 def _get_typing_origin(t: type) -> type:
     """
@@ -589,7 +641,5 @@ def _create_type_mapping():
 
             ImagePreset._TYPE_MAPPING[type_key] = type_value  # noqa
 
-    assert all(type_key in ImagePreset._TYPE_MAPPING for type_key in ImagePreset._DEFAULT)  # noqa
-
 
 _create_type_mapping()

novelai_api/_high_level.py

@@ -489,7 +489,7 @@ async def generate_image(
                 ImageModel.Anime_Curated,
                 ImageModel.Furry,
                 ImageModel.Inpainting_Anime_Full,
-                ImageModel.Inainting_Anime_Curated,
+                ImageModel.Inpainting_Anime_Curated,
                 ImageModel.Inpainting_Furry,
             ):
                 prompt = f"masterpiece, best quality, {prompt}"