import copy
import enum
import json
import math
import os
import random
import sys
import typing
from pathlib import Path
from typing import Any, Dict, List, 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
[docs]class ImageModel(enum.Enum):
"""
Image model for low_level.suggest_tags() and low_level.generate_image()
"""
Anime_Curated = "safe-diffusion"
Anime_Full = "nai-diffusion"
Furry = "nai-diffusion-furry"
Inpainting_Anime_Curated = "safe-diffusion-inpainting"
Inpainting_Anime_Full = "nai-diffusion-inpainting"
Inpainting_Furry = "furry-diffusion-inpainting"
Anime_v2 = "nai-diffusion-2"
Anime_v3 = "nai-diffusion-3"
Inpainting_Anime_v3 = "nai-diffusion-3-inpainting"
Furry_v3 = "nai-diffusion-furry-3"
Inpainting_Furry_v3 = "nai-diffusion-furry-3-inpainting"
Anime_v4_preview = "nai-diffusion-4-curated-preview"
[docs]class ControlNetModel(enum.Enum):
"""
ControlNet Model for ImagePreset.controlnet_model and low_level.generate_controlnet_mask()
"""
Palette_Swap = "hed"
Form_Lock = "midas"
Scribbler = "fake_scribble"
Building_Control = "mlsd"
Landscaper = "uniformer"
[docs]class ImageResolution(enum.Enum):
"""
Image resolution for ImagePreset.resolution
"""
Wallpaper_Portrait = (1088, 1920)
Wallpaper_Landscape = (1920, 1088)
# v1
Small_Portrait = (384, 640)
Small_Landscape = (640, 384)
Small_Square = (512, 512)
Normal_Portrait = (512, 768)
Normal_Landscape = (768, 512)
Normal_Square = (640, 640)
Large_Portrait = (512, 1024)
Large_Landscape = (1024, 512)
Large_Square = (1024, 1024)
# v2
Small_Portrait_v2 = (512, 768)
Small_Landscape_v2 = (768, 512)
Small_Square_v2 = (640, 640)
Normal_Portrait_v2 = (832, 1216)
Normal_Landscape_v2 = (1216, 832)
Normal_Square_v2 = (1024, 1024)
Large_Portrait_v2 = (1024, 1536)
Large_Landscape_v2 = (1536, 1024)
Large_Square_v2 = (1472, 1472)
# v3
Small_Portrait_v3 = (512, 768)
Small_Landscape_v3 = (768, 512)
Small_Square_v3 = (640, 640)
Normal_Portrait_v3 = (832, 1216)
Normal_Landscape_v3 = (1216, 832)
Normal_Square_v3 = (1024, 1024)
Large_Portrait_v3 = (1024, 1536)
Large_Landscape_v3 = (1536, 1024)
Large_Square_v3 = (1472, 1472)
# v4
Small_Portrait_v4 = (512, 768)
Small_Landscape_v4 = (768, 512)
Small_Square_v4 = (640, 640)
Normal_Portrait_v4 = (832, 1216)
Normal_Landscape_v4 = (1216, 832)
Normal_Square_v4 = (1024, 1024)
Large_Portrait_v4 = (1024, 1536)
Large_Landscape_v4 = (1536, 1024)
Large_Square_v4 = (1472, 1472)
[docs]class ImageSampler(enum.Enum):
"""
Sampler for ImagePreset.sampler
"""
k_lms = "k_lms"
k_euler = "k_euler"
k_euler_ancestral = "k_euler_ancestral"
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"
k_dpmpp_2m = "k_dpmpp_2m"
k_dpmpp_2m_sde = "k_dpmpp_2m_sde"
k_dpmpp_2s_ancestral = "k_dpmpp_2s_ancestral"
k_dpmpp_sde = "k_dpmpp_sde"
k_dpm_2 = "k_dpm_2"
k_dpm_2_ancestral = "k_dpm_2_ancestral"
k_dpm_adaptive = "k_dpm_adaptive" # doesn't work
k_dpm_fast = "k_dpm_fast" # doesn't work
[docs]class UCPreset(enum.Enum):
"""
Default UC preset for ImagePreset.uc_preset
"""
Preset_Low_Quality_Bad_Anatomy = 0
Preset_Low_Quality = 1
Preset_Bad_Anatomy = 2
Preset_None = 3
Preset_Heavy = 4
Preset_Light = 5
[docs]class ImageGenerationType(enum.Enum):
"""
Image generation type for low_level.generate_image
"""
NORMAL = "generate"
IMG2IMG = "img2img"
INPAINTING = "infill"
[docs]class ImagePreset:
_UC_Presets = {
# v1
ImageModel.Anime_Curated: {
UCPreset.Preset_Low_Quality_Bad_Anatomy: "nsfw, lowres, bad anatomy, bad hands, text, error, "
"missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, "
"jpeg artifacts, signature, watermark, username, blurry",
UCPreset.Preset_Bad_Anatomy: None,
UCPreset.Preset_Low_Quality: "nsfw, lowres, text, cropped, worst quality, low quality, normal quality, "
"jpeg artifacts, signature, watermark, twitter username, blurry",
UCPreset.Preset_None: "lowres",
},
ImageModel.Anime_Full: {
UCPreset.Preset_Low_Quality_Bad_Anatomy: "nsfw, lowres, bad anatomy, bad hands, text, error, "
"missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, "
"jpeg artifacts, signature, watermark, username, blurry",
UCPreset.Preset_Bad_Anatomy: None,
UCPreset.Preset_Low_Quality: "nsfw, lowres, text, cropped, worst quality, low quality, normal quality, "
"jpeg artifacts, signature, watermark, twitter username, blurry",
UCPreset.Preset_None: "lowres",
},
ImageModel.Furry: {
UCPreset.Preset_Low_Quality_Bad_Anatomy: None,
UCPreset.Preset_Low_Quality: "nsfw, worst quality, low quality, what has science done, what, "
"nightmare fuel, eldritch horror, where is your god now, why",
UCPreset.Preset_Bad_Anatomy: "{worst quality}, low quality, distracting watermark, [nightmare fuel], "
"{{unfinished}}, deformed, outline, pattern, simple background",
UCPreset.Preset_None: "low res",
},
# v2
ImageModel.Anime_v2: {
UCPreset.Preset_Heavy: "nsfw, lowres, bad, text, error, missing, extra, fewer, cropped, jpeg artifacts, "
"worst quality, bad quality, watermark, displeasing, unfinished, chromatic aberration, scan, "
"scan artifacts",
UCPreset.Preset_Light: "nsfw, lowres, jpeg artifacts, worst quality, watermark, blurry, very displeasing",
UCPreset.Preset_None: "lowres",
},
# v3
ImageModel.Anime_v3: {
UCPreset.Preset_Heavy: "nsfw, lowres, {bad}, error, fewer, extra, missing, worst quality, jpeg artifacts, "
"bad quality, watermark, unfinished, displeasing, chromatic aberration, signature, extra digits, "
"artistic error, username, scan, [abstract]",
UCPreset.Preset_Light: "nsfw, lowres, jpeg artifacts, worst quality, watermark, blurry, very displeasing",
UCPreset.Preset_None: "lowres",
},
ImageModel.Furry_v3: {
UCPreset.Preset_Heavy: "nsfw, {{worst quality}}, [displeasing], {unusual pupils}, guide lines, "
"{{unfinished}}, {bad}, url, artist name, {{tall image}}, mosaic, {sketch page}, comic panel, "
"impact (font), [dated], {logo}, ych, {what}, {where is your god now}, {distorted text}, repeated text, "
"{floating head}, {1994}, {widescreen}, absolutely everyone, sequence, {compression artifacts}, "
"hard translated, {cropped}, {commissioner name}, unknown text, high contrast",
UCPreset.Preset_Light: "{worst quality}, guide lines, unfinished, bad, url, tall image, widescreen, "
"compression artifacts, unknown text",
UCPreset.Preset_None: "lowres",
},
# v4
ImageModel.Anime_v4_preview: {
UCPreset.Preset_Heavy: "blurry, lowres, error, film grain, scan artifacts, worst quality, bad quality, "
"jpeg artifacts, very displeasing, chromatic aberration, logo, dated, signature, multiple views, "
"gigantic breasts",
UCPreset.Preset_Light: "blurry, lowres, error, worst quality, bad quality, jpeg artifacts, "
"very displeasing, logo, dated, signature",
UCPreset.Preset_None: "lowres",
},
}
# inpainting presets are the same as the normal ones
_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]
_CONTROLNET_MODELS = {
ControlNetModel.Palette_Swap: "hed",
ControlNetModel.Form_Lock: "depth",
ControlNetModel.Scribbler: "scribble",
ControlNetModel.Building_Control: "mlsd",
ControlNetModel.Landscaper: "seg",
}
_TYPE_MAPPING = {
"legacy": bool,
# rest is populated in at the bottom of the file
}
# 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
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
#: 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
#: reference_image for multi-vibe transfer
reference_image_multiple: List[str]
#: reference_information_extracted for multi-vibe transfer
reference_information_extracted_multiple: List[float]
#: reference_strength for multi-vibe transfer
reference_strength_multiple: List[float]
#: https://blog.novelai.net/summer-sampler-update-en-3a34eb32b613
variety_plus: bool
#: Whether the AI should strictly follow the positions of the characters or have some freedom
use_coords: bool
#: https://docs.novelai.net/image/multiplecharacters.html#multi-character-prompting
#: See examples/generate_image_v4.py for the format
characters: List[Dict[str, str]]
#: Use the old behavior of prompt separation at the 75 tokens mark (can cut words in half)
legacy_v3_extend: bool
#: Revision of the default arguments
params_version: int
#: Use the old behavior of noise scheduling with the k_euler_ancestral sampler
deliberate_euler_ancestral_bug: bool
#: ???
prefer_brownian: 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
[docs] @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)
[docs] 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))
[docs] @expand_kwargs(_TYPE_MAPPING.keys(), _TYPE_MAPPING.values())
def __init__(self, **kwargs):
"""
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)
[docs] @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")
[docs] @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")
[docs] @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")
[docs] @classmethod
def from_v3_furry_config(cls):
"""
Create a new ImagePreset with the default settings from the v3 furry config
"""
return cls.from_file(Path(__file__).parent / "image_presets" / "presets_v3" / "default_furry.preset")
[docs] @classmethod
def from_v4_config(cls):
"""
Create a new ImagePreset with the default settings from the v4 config
"""
return cls.from_file(Path(__file__).parent / "image_presets" / "presets_v4" / "default.preset")
[docs] @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()
elif model in (ImageModel.Furry_v3, ImageModel.Inpainting_Furry_v3):
return cls.from_v3_furry_config()
elif model in (ImageModel.Anime_v4_preview,):
return cls.from_v4_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)}'")
self._settings[key] = value
def __getitem__(self, key: str):
return self._settings[key]
def __delitem__(self, key):
if key in self._DEFAULT:
raise ValueError(f"'{key}' is a default setting, set it instead of deleting")
del self._settings[key]
def __contains__(self, key: str):
return key in self._settings.keys()
[docs] def update(self, values: Optional[Dict[str, Any]] = None, **kwargs) -> "ImagePreset":
"""
Update the settings stored in the preset. Works like dict.update()
"""
if values is not None:
for k, v in values.items():
self[k] = v
for k, v in kwargs.items():
self[k] = v
return self
[docs] def copy(self) -> "ImagePreset":
"""
Create a new ImagePreset instance from the current one
"""
return ImagePreset(**self._settings)
# give dot access capabilities to the object
def __setattr__(self, key, value):
if key in self._TYPE_MAPPING:
self[key] = value
else:
object.__setattr__(self, key, value)
def __getattr__(self, key):
if key in self._TYPE_MAPPING:
return self[key]
return object.__getattribute__(self, key)
def __delattr__(self, name):
if name in self._TYPE_MAPPING:
del self[name]
else:
object.__delattr__(self, name)
[docs] def to_settings(self, model: ImageModel) -> Dict[str, Any]:
"""
Return the values stored in the preset, for a generate_image function
:param model: Image model to get the settings of
"""
settings = copy.deepcopy(self._settings)
# size
resolution: Union[ImageResolution, Tuple[int, int]] = settings.pop("resolution")
if isinstance(resolution, ImageResolution):
resolution: Tuple[int, int] = resolution.value
settings["width"], settings["height"] = resolution
# seed 0 = random seed for the backend, but it is not set in metadata, so we set it ourself to be safe
# the seed of the ith image is seed + i, so we reserve space for them (makes valid images with invalid metadata)
seed = settings.pop("seed")
if seed == 0:
seed = random.randint(1, 0xFFFFFFFF - settings["n_samples"] + 1)
self.last_seed = seed
settings["seed"] = seed
settings["extra_noise_seed"] = seed
# UC
uc_preset: Union[UCPreset, None] = settings.pop("uc_preset")
if uc_preset is None:
default_uc = ""
else:
default_uc = self._UC_Presets[model].get(uc_preset, None)
if default_uc is None:
raise ValueError(f"UC preset '{uc_preset.name}' is not valid for model '{model.value}'")
uc: str = settings.pop("uc")
combined_uc = f"{default_uc}, {uc}" if default_uc and uc else default_uc if default_uc else uc
settings["negative_prompt"] = combined_uc
# sampler
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)
settings["sm_dyn"] = settings.pop("smea_dyn", False)
controlnet_model: Optional[ControlNetModel] = settings.pop("controlnet_model", None)
if controlnet_model is not None:
settings["controlnet_model"] = self._CONTROLNET_MODELS[controlnet_model]
settings["dynamic_thresholding"] = settings.pop("decrisper")
settings["skip_cfg_above_sigma"] = 19 if settings.pop("variety_plus", False) else None
# character prompts
if model in (ImageModel.Anime_v4_preview,):
settings["v4_prompt"] = {
# base_caption is set later, in generate_image
"caption": {"base_caption": None, "char_captions": []},
"use_coords": self.use_coords,
"use_order": True,
}
settings["v4_negative_prompt"] = {"caption": {"base_caption": combined_uc, "char_captions": []}}
characters = settings.pop("characters", [])
if not isinstance(characters, list):
raise ValueError("characters must be a list of dictionaries")
settings["characterPrompts"] = []
for i, character in enumerate(characters):
if not isinstance(character, dict):
raise ValueError(f"character #{i} must be a dictionary")
if "prompt" not in character:
raise ValueError(f"character #{i} must have at least a 'prompt' key")
prompt = character["prompt"]
if not isinstance(prompt, str):
raise ValueError(f"character #{i} prompt must be a string")
negative = character.get("uc", "")
character_position = character.get("position", "") or "C3"
if (
len(character_position) != 2
or character_position[0] not in "ABCDE"
or character_position[1] not in "12345"
):
raise ValueError(f'character #{i} position must be a valid position ("", or "A1" to "E5")')
pos = {
"x": round(0.5 + 0.2 * (ord(character_position[0]) - ord("C")), 1),
"y": round(0.5 + 0.2 * (ord(character_position[1]) - ord("3")), 1),
}
settings["characterPrompts"].append({"center": pos, "prompt": prompt, "uc": negative})
settings["v4_prompt"]["caption"]["char_captions"].append({"centers": [pos], "char_caption": prompt})
settings["v4_negative_prompt"]["caption"]["char_captions"].append(
{"centers": [pos], "char_caption": negative}
)
# special arguments kept for metadata purposes (no effect on result)
settings["qualityToggle"] = settings.pop("quality_toggle")
if uc_preset is not None:
settings["ucPreset"] = uc_preset.value
return settings
[docs] def get_max_n_samples(self):
"""
Get the allowed max value of ImagePreset.n_samples using current preset values
"""
resolution: Union[ImageResolution, Tuple[int, int]] = self._settings["resolution"]
if isinstance(resolution, ImageResolution):
resolution: Tuple[int, int] = resolution.value
w, h = resolution
if w * h <= 512 * 704:
return 8
if w * h <= 640 * 640:
return 6
if w * h <= 512 * 2560:
return 4
if w * h <= 1024 * 1536:
return 2
if w * h <= 1024 * 3072:
return 1
return 0
[docs] def calculate_cost(
self, is_opus: bool, version: int = 1, generation_type: ImageGenerationType = ImageGenerationType.NORMAL
):
"""
Calculate the cost (in Anlas) of generating with the current configuration
:param is_opus: Is the subscription tier Opus ? Account for free generations if so
:param version: Version of the model to use (1, 2, 3)
:param generation_type: Type of generation to do (img2img, txt2img, etc.)
"""
steps: int = self._settings["steps"]
n_samples: int = max(1, self._settings["n_samples"])
smea = self._settings["smea"]
smea_dyn = self._settings["smea_dyn"]
sampler: ImageSampler = self._settings["sampler"]
uncond_scale: float = self._settings.get("uncond_scale", 1.0)
strength: float = self._settings.get("strength", 1.0) if generation_type == ImageGenerationType.IMG2IMG else 1.0
resolution: Union[ImageResolution, Tuple[int, int]] = self._settings["resolution"]
if isinstance(resolution, ImageResolution):
resolution: Tuple[int, int] = resolution.value
w, h = resolution
r = w * h
if r < 65536:
r = 65536
if version == 3:
smea_factor = 1.0 if not smea else 1.2 if not smea_dyn else 1.4
per_sample = math.ceil(2951823174884865e-21 * r + 5.753298233447344e-7 * r * steps) * smea_factor
else:
if r <= 1024 * 1024 and sampler in (
ImageSampler.plms,
ImageSampler.ddim,
ImageSampler.k_euler,
ImageSampler.k_euler_ancestral,
ImageSampler.k_lms,
):
per_sample = (
(15.266497014243718 * math.exp(r / 1024 / 1024 * 0.6326248927474729) - 15.225164493059737)
* steps
/ 28
)
else:
index = math.ceil(w / 64) * math.ceil(h / 64) - 1
if sampler is ImageSampler.nai_smea_dyn or (smea and smea_dyn):
per_step, fixed = SMEA_DYN_COSTS[index]
elif sampler is ImageSampler.nai_smea or smea:
per_step, fixed = SMEA_COSTS[index]
elif sampler is ImageSampler.ddim:
per_step, fixed = DDIM_COSTS[index]
else:
per_step, fixed = NAI_COSTS[index]
per_sample = per_step * steps + fixed
per_sample = max(math.ceil(per_sample * strength), 2)
if version != 1 and uncond_scale != 1.0:
per_sample = math.ceil(per_sample * uncond_scale)
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))
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
_create_type_mapping()