From 0f62c89c1fde1f50fb422a2886d1ed0320f22f4f Mon Sep 17 00:00:00 2001 From: Aedial Date: Tue, 5 Dec 2023 00:51:37 +0100 Subject: [PATCH] [API] Add automatic TYPE_MAPPING to ImagePreset and fix ddim on v3 ddim is now converted into ddim_v3 on Anime_v3 --- novelai_api/ImagePreset.py | 195 +++++++++++++++++++++++-------------- 1 file changed, 124 insertions(+), 71 deletions(-) diff --git a/novelai_api/ImagePreset.py b/novelai_api/ImagePreset.py index bdacc8b..3776801 100644 --- a/novelai_api/ImagePreset.py +++ b/novelai_api/ImagePreset.py @@ -4,7 +4,9 @@ import math import os import random -from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union +import sys +import typing +from typing import Any, Dict, Optional, Tuple, Union from novelai_api.ImagePreset_CostTables import DDIM_COSTS, NAI_COSTS, SMEA_COSTS, SMEA_DYN_COSTS from novelai_api.python_utils import NoneType, expand_kwargs @@ -100,6 +102,7 @@ class ImageSampler(enum.Enum): k_heun = "k_heun" plms = "plms" # doesn't work ddim = "ddim" + ddim_v3 = "ddim_v3" # for v3 nai_smea = "nai_smea" # doesn't work nai_smea_dyn = "nai_smea_dyn" @@ -200,79 +203,56 @@ class ImagePreset: _TYPE_MAPPING = { "legacy": bool, - "quality_toggle": bool, - "resolution": (ImageResolution, tuple), - "uc_preset": (UCPreset, NoneType), - "n_samples": int, - "seed": int, - "sampler": ImageSampler, - "noise": (int, float), - "strength": (int, float), - "scale": (int, float), - "uncond_scale": (int, float), - "steps": int, - "uc": str, - "smea": bool, - "smea_dyn": bool, - "image": str, - "controlnet_condition": str, - "controlnet_model": ControlNetModel, - "controlnet_strength": (int, float), - "decrisper": bool, - "add_original_image": bool, - "mask": str, - "cfg_rescale": float, - "noise_schedule": str, + # rest is populated in at the bottom of the file } # type completion for __setitem__ and __getitem__ - if TYPE_CHECKING: - #: https://docs.novelai.net/image/qualitytags.html - quality_toggle: 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 - uc_preset: Union[UCPreset, None] - #: Number of images to return - n_samples: int - #: Random seed to use for the image. The ith image has seed + i for seed - seed: int - #: https://docs.novelai.net/image/sampling.html - sampler: ImageSampler - #: https://docs.novelai.net/image/strengthnoise.html - noise: float - #: https://docs.novelai.net/image/strengthnoise.html - strength: float - #: https://docs.novelai.net/image/stepsguidance.html (scale is called Prompt Guidance) - scale: float - #: TODO - uncond_scale: float - #: https://docs.novelai.net/image/stepsguidance.html - steps: int - #: https://docs.novelai.net/image/undesiredcontent.html - uc: str - #: Enable SMEA for any sampler (makes Large+ generations manageable) - smea: bool - #: Enable SMEA DYN for any sampler if SMEA is enabled (best for Large+, but not Wallpaper resolutions) - smea_dyn: bool - #: b64-encoded png image for img2img - image: str - #: Controlnet mask gotten by the generate_controlnet_mask method - controlnet_condition: str - #: Model to use for the controlnet - controlnet_model: ControlNetModel - #: Influence of the chosen controlnet on the image - controlnet_strength: float - #: Reduce the deepfrying effects of high scale (https://twitter.com/Birchlabs/status/1582165379832348672) - decrisper: bool - #: Prevent seams along the edges of the mask, but may change the image slightly - add_original_image: bool - #: Mask for inpainting (b64-encoded black and white png image, white is the inpainting area) - mask: str - #: https://docs.novelai.net/image/stepsguidance.html#prompt-guidance-rescale - cfg_rescale: float - #: ??? (TODO: use an enum ? - valid values: native, karras, exponential, polyexponential) - noise_schedule: str + #: https://docs.novelai.net/image/qualitytags.html + quality_toggle: 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 + uc_preset: Union[UCPreset, None] + #: Number of images to return + n_samples: int + #: Random seed to use for the image. The ith image has seed + i for seed + seed: int + #: https://docs.novelai.net/image/sampling.html + sampler: ImageSampler + #: https://docs.novelai.net/image/strengthnoise.html + noise: float + #: https://docs.novelai.net/image/strengthnoise.html + strength: float + #: https://docs.novelai.net/image/stepsguidance.html (scale is called Prompt Guidance) + scale: float + #: TODO + uncond_scale: float + #: https://docs.novelai.net/image/stepsguidance.html + steps: int + #: https://docs.novelai.net/image/undesiredcontent.html + uc: str + #: Enable SMEA for any sampler (makes Large+ generations manageable) + smea: bool + #: Enable SMEA DYN for any sampler if SMEA is enabled (best for Large+, but not Wallpaper resolutions) + smea_dyn: bool + #: b64-encoded png image for img2img + image: str + #: Controlnet mask gotten by the generate_controlnet_mask method + controlnet_condition: str + #: Model to use for the controlnet + controlnet_model: ControlNetModel + #: Influence of the chosen controlnet on the image + controlnet_strength: float + #: Reduce the deepfrying effects of high scale (https://twitter.com/Birchlabs/status/1582165379832348672) + decrisper: bool + #: Prevent seams along the edges of the mask, but may change the image slightly + add_original_image: bool + #: Mask for inpainting (b64-encoded black and white png image, white is the inpainting area) + mask: str + #: https://docs.novelai.net/image/stepsguidance.html#prompt-guidance-rescale + cfg_rescale: float + #: ??? (TODO: use an enum ? - valid values: native, karras, exponential, polyexponential) + noise_schedule: str _DEFAULT = { "legacy": False, @@ -405,6 +385,9 @@ def to_settings(self, model: ImageModel) -> Dict[str, Any]: settings["negative_prompt"] = combined_uc sampler: ImageSampler = settings.pop("sampler") + if sampler is ImageSampler.ddim and model in (ImageModel.Anime_v3,): + sampler = ImageSampler.ddim_v3 + settings["sampler"] = sampler.value settings["sm"] = settings.pop("smea", False) @@ -540,3 +523,73 @@ def to_file(self, path: Union[str, bytes, os.PathLike, int]): with open(path, "w", encoding="utf-8") as f: f.write(json.dumps(self._settings)) + + +def _get_typing_origin(t: type) -> type: + """ + Get the typing origin of a type + + :param t: Type to get the origin of + """ + + if sys.version_info < (3, 8): # 3.7 + origin = getattr(t, "__origin__", None) + assert origin is not None # should never happen for 3.7 + + return origin + + return typing.get_origin(t) + + +def _get_typing_args(t: type) -> Tuple[type, ...]: + """ + Get the typing arguments of a type + + :param t: Type to get the arguments of + """ + + if sys.version_info < (3, 8): # 3.7 + args = getattr(t, "__args__", None) + assert args is not None # should never happen for 3.7 + + return args + + return typing.get_args(t) + + +def _get_recursive_type(t: type, depth: int = 1) -> Union[type, Tuple[type, ...]]: + if t is None: + return NoneType + + if t.__module__ == "typing": + if _get_typing_origin(t) is Union: + if depth == 0: + raise ValueError("Union types are not supported past depth 1") + + return tuple(_get_recursive_type(x, depth - 1) for x in _get_typing_args(t)) + + return _get_typing_origin(t) + + return t + + +def _create_type_mapping(): + """ + Create the type mapping for the ImagePreset class + """ + + non_mapping_keys = ["last_seed"] + + for type_key, type_value in ImagePreset.__annotations__.items(): + if not type_key.startswith("_") and type_key != type_key.upper() and type_key not in non_mapping_keys: + if type_value is float: + type_value = (int, float) + else: + type_value = _get_recursive_type(type_value) + + 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()