novelai-python 0.7.5__tar.gz → 0.7.6__tar.gz

{novelai_python-0.7.5 → novelai_python-0.7.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: novelai-python
-Version: 0.7.5
+Version: 0.7.6
 Summary: NovelAI Python Binding With Pydantic
 Author-Email: sudoskys <coldlando@hotmail.com>
 License: Apache License 2.0

{novelai_python-0.7.5 → novelai_python-0.7.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.7.5"
+version = "0.7.6"
 description = "NovelAI Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "coldlando@hotmail.com" },

{novelai_python-0.7.5 → novelai_python-0.7.6}/src/novelai_python/sdk/ai/_enum.py

@@ -92,6 +92,9 @@ class Resolution(Enum):
 
 class Model(Enum):
     NAI_DIFFUSION_4_CURATED_PREVIEW = "nai-diffusion-4-curated-preview"
+    NAI_DIFFUSION_4_FULL = "nai-diffusion-4-full"
+    NAI_DIFFUSION_4_FULL_INPAINTING = "nai-diffusion-4-full-inpainting"
+    NAI_DIFFUSION_4_CURATED_INPAINTING = "nai-diffusion-4-curated-inpainting"
 
     NAI_DIFFUSION_3 = "nai-diffusion-3"
     NAI_DIFFUSION_3_INPAINTING = "nai-diffusion-3-inpainting"
@@ -151,6 +154,8 @@ PROMOTION = {
     "Stable Diffusion XL C8704949": Model.NAI_DIFFUSION_FURRY_3,
     "Stable Diffusion XL 37C2B166": Model.NAI_DIFFUSION_FURRY_3,
     "Stable Diffusion XL 9CC2F394": Model.NAI_DIFFUSION_FURRY_3,
+    'NovelAI Diffusion V4 4F49EC75': Model.NAI_DIFFUSION_4_FULL,
+    'NovelAI Diffusion V4 CA4B7203': Model.NAI_DIFFUSION_4_FULL,
 }
 
 ModelTypeAlias = Optional[Union[Model, str]]
@@ -171,6 +176,7 @@ class SupportCondition:
     v4Prompts: bool
     smea: bool
     text: bool
+    dynamicThresholding: bool
 
 
 def get_supported_params(model: Model):
@@ -201,7 +207,8 @@ def get_supported_params(model: Model):
             characterPrompts=False,
             v4Prompts=False,
             smea=True,
-            text=False
+            text=False,
+            dynamicThresholding=True
         )
     if model in [Model.NAI_DIFFUSION_2]:
         return SupportCondition(
@@ -215,7 +222,8 @@ def get_supported_params(model: Model):
             characterPrompts=False,
             v4Prompts=False,
             smea=True,
-            text=False
+            text=False,
+            dynamicThresholding=True
         )
     if model in [
         Model.NAI_DIFFUSION_XL,
@@ -235,21 +243,29 @@ def get_supported_params(model: Model):
             characterPrompts=False,
             v4Prompts=False,
             smea=True,
-            text=False
+            text=False,
+            dynamicThresholding=True
         )
-    if model in [Model.CUSTOM, Model.NAI_DIFFUSION_4_CURATED_PREVIEW]:
+    if model in [
+        Model.NAI_DIFFUSION_4_CURATED_PREVIEW,
+        Model.NAI_DIFFUSION_4_CURATED_INPAINTING,
+        Model.CUSTOM,
+        Model.NAI_DIFFUSION_4_FULL,
+        Model.NAI_DIFFUSION_4_FULL_INPAINTING,
+    ]:
         return SupportCondition(
             controlnet=False,
             vibetransfer=False,
-            scaleMax=10,
+            scaleMax=20,
             negativePromptGuidance=True,
             noiseSchedule=True,
             inpainting=True,
-            cfgDelay=False,
+            cfgDelay=True,
             characterPrompts=True,
             v4Prompts=True,
             smea=False,
-            text=True
+            text=True,
+            dynamicThresholding=False
         )
     return SupportCondition(
         controlnet=False,
@@ -262,7 +278,8 @@ def get_supported_params(model: Model):
         characterPrompts=False,
         v4Prompts=False,
         smea=False,
-        text=False
+        text=False,
+        dynamicThresholding=False
     )
 
 
@@ -272,13 +289,31 @@ class Modifier(object):
     suffix: str
 
 
+def find_model_by_hashcode(hashcode: str) -> ModelTypeAlias:
+    if "NovelAI Diffusion V4" in hashcode:
+        return PROMOTION.get(hashcode, Model.NAI_DIFFUSION_4_CURATED_PREVIEW)
+    return PROMOTION.get(hashcode, None)
+
+
 def get_modifiers(model: Model) -> Modifier:
     """
     Get modifiers from model
     :param model:
     :return:
     """
-    if model in [Model.CUSTOM, Model.NAI_DIFFUSION_4_CURATED_PREVIEW]:
+    if model in [
+        Model.CUSTOM,
+        Model.NAI_DIFFUSION_4_FULL,
+        Model.NAI_DIFFUSION_4_FULL_INPAINTING,
+    ]:
+        return Modifier(
+            qualityTags="",
+            suffix=", no text, best quality, very aesthetic, absurdres"
+        )
+    if model in [
+        Model.NAI_DIFFUSION_4_CURATED_PREVIEW,
+        Model.NAI_DIFFUSION_4_CURATED_INPAINTING,
+    ]:
         return Modifier(
             qualityTags="",
             suffix=", rating:general, best quality, very aesthetic, absurdres"
@@ -310,7 +345,37 @@ def get_modifiers(model: Model) -> Modifier:
     )
 
 
-def get_supported_noise_schedule(sample_type: Sampler) -> List[NoiseSchedule]:
+def get_supported_noise_schedule(sample_type: Sampler, model: Model) -> List[NoiseSchedule]:
+    """
+    Get supported noise schedule for a given sample type and model
+    :param sample_type: Sampler
+    :param model: Model
+    :return: List[NoiseSchedule]
+    """
+    noise_schedule = get_sampler_supported_noise_schedule(sample_type)
+    schedules = filter_by_model_noise_schedule(model, noise_schedule)
+    return schedules
+
+
+def filter_by_model_noise_schedule(model: Model, noise_schedule: List[NoiseSchedule]) -> List[NoiseSchedule]:
+    """
+    Filter noise schedule by model
+    :param model:
+    :param noise_schedule:
+    :return:
+    """
+    # 只有Novelai4系列不支持Native
+    if model in [
+        Model.NAI_DIFFUSION_4_CURATED_PREVIEW,
+        Model.NAI_DIFFUSION_4_FULL,
+        Model.NAI_DIFFUSION_4_FULL_INPAINTING,
+        Model.NAI_DIFFUSION_4_CURATED_INPAINTING
+    ]:
+        return [schedule for schedule in noise_schedule if schedule != NoiseSchedule.NATIVE]
+    return noise_schedule
+
+
+def get_sampler_supported_noise_schedule(sample_type: Sampler) -> List[NoiseSchedule]:
     """
     Get supported noise schedule for a given sample type
     :param sample_type: Sampler
@@ -384,6 +449,9 @@ def get_model_group(model: ModelTypeAlias) -> ModelGroups:
         "nai-diffusion-furry-3": ModelGroups.STABLE_DIFFUSION_XL_FURRY,
         "nai-diffusion-furry-3-inpainting": ModelGroups.STABLE_DIFFUSION_XL_FURRY,
         "nai-diffusion-4-curated-preview": ModelGroups.V4,
+        'nai-diffusion-4-full': ModelGroups.V4,
+        'nai-diffusion-4-full-inpainting': ModelGroups.V4,
+        'nai-diffusion-4-curated-inpainting': ModelGroups.V4
     }
     return mapping.get(model, ModelGroups.STABLE_DIFFUSION)
 
@@ -404,67 +472,152 @@ def get_uc_preset(model: ModelTypeAlias) -> List[UcPrompt]:
         Model.SAFE_DIFFUSION_INPAINTING,
     ]:
         prompts = [
-            UcPrompt(category="heavy", name="lowQualityPlusBadAnatomy",
-                     text="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"),
-            UcPrompt(category="light", name="lowQuality",
-                     text="lowres, text, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry"),
-            UcPrompt(category="none", name="none", text="lowres"),
+            UcPrompt(
+                category="heavy",
+                name="lowQualityPlusBadAnatomy",
+                text="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"
+            ),
+            UcPrompt(
+                category="light",
+                name="lowQuality",
+                text="lowres, text, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry"
+            ),
+            UcPrompt(
+                category="none",
+                name="none",
+                text="lowres"
+            ),
         ]
     elif model in [
         Model.NAI_DIFFUSION_FURRY,
         Model.FURRY_DIFFUSION_INPAINTING,
     ]:
         prompts = [
-            UcPrompt(category="light", name="lowQuality",
-                     text="worst quality, low quality, what has science done, what, nightmare fuel, eldritch horror, where is your god now, why"),
-            UcPrompt(category="heavy", name="badAnatomy",
-                     text="{worst quality}, low quality, distracting watermark, [nightmare fuel], {{unfinished}}, deformed, outline, pattern, simple background"),
-            UcPrompt(category="none", name="none", text="low res"),
+            UcPrompt(
+                category="light",
+                name="lowQuality",
+                text="worst quality, low quality, what has science done, what, nightmare fuel, eldritch horror, where is your god now, why"
+            ),
+            UcPrompt(
+                category="heavy",
+                name="badAnatomy",
+                text="{worst quality}, low quality, distracting watermark, [nightmare fuel], {{unfinished}}, deformed, outline, pattern, simple background"
+            ),
+            UcPrompt(
+                category="none",
+                name="none",
+                text="low res"
+            ),
         ]
     elif model in [
         Model.NAI_DIFFUSION_2,
     ]:
         prompts = [
-            UcPrompt(category="heavy", name="heavy",
-                     text="lowres, bad, text, error, missing, extra, fewer, cropped, jpeg artifacts, worst quality, bad quality, watermark, displeasing, unfinished, chromatic aberration, scan, scan artifacts"),
-            UcPrompt(category="light", name="light",
-                     text="lowres, jpeg artifacts, worst quality, watermark, blurry, very displeasing"),
-            UcPrompt(category="none", name="none", text="lowres"),
+            UcPrompt(
+                category="heavy",
+                name="heavy",
+                text="lowres, bad, text, error, missing, extra, fewer, cropped, jpeg artifacts, worst quality, bad quality, watermark, displeasing, unfinished, chromatic aberration, scan, scan artifacts"
+            ),
+            UcPrompt(
+                category="light",
+                name="light",
+                text="lowres, jpeg artifacts, worst quality, watermark, blurry, very displeasing"
+            ),
+            UcPrompt(
+                category="none",
+                name="none",
+                text="lowres"
+            ),
         ]
     elif model in [
         Model.NAI_DIFFUSION_3,
         Model.NAI_DIFFUSION_3_INPAINTING,
     ]:
         prompts = [
-            UcPrompt(category="heavy", name="heavy",
-                     text="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]"),
-            UcPrompt(category="light", name="light",
-                     text="lowres, jpeg artifacts, worst quality, watermark, blurry, very displeasing"),
-            UcPrompt(category="human", name="humanFocus",
-                     text="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], bad anatomy, bad hands, @_@, mismatched pupils, heart-shaped pupils, glowing eyes"),
-            UcPrompt(category="none", name="none", text="lowres"),
+            UcPrompt(
+                category="heavy",
+                name="heavy",
+                text="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]"
+            ),
+            UcPrompt(
+                category="light",
+                name="light",
+                text="lowres, jpeg artifacts, worst quality, watermark, blurry, very displeasing"
+            ),
+            UcPrompt(
+                category="human",
+                name="humanFocus",
+                text="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], bad anatomy, bad hands, @_@, mismatched pupils, heart-shaped pupils, glowing eyes"
+            ),
+            UcPrompt(
+                category="none",
+                name="none",
+                text="lowres"
+            ),
         ]
     elif model in [
         Model.NAI_DIFFUSION_FURRY_3,
         Model.NAI_DIFFUSION_FURRY_3_INPAINTING,
     ]:
         prompts = [
-            UcPrompt(category="heavy", name="heavy",
-                     text="{{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"),
-            UcPrompt(category="light", name="light",
-                     text="{worst quality}, guide lines, unfinished, bad, url, tall image, widescreen, compression artifacts, unknown text"),
-            UcPrompt(category="none", name="none", text="lowres"),
+            UcPrompt(
+                category="heavy",
+                name="heavy",
+                text="{{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"
+            ),
+            UcPrompt(
+                category="light",
+                name="light",
+                text="{worst quality}, guide lines, unfinished, bad, url, tall image, widescreen, compression artifacts, unknown text"
+            ),
+            UcPrompt(
+                category="none",
+                name="none",
+                text="lowres"
+            ),
         ]
     elif model in [
-        Model.CUSTOM,
         Model.NAI_DIFFUSION_4_CURATED_PREVIEW,
+        Model.NAI_DIFFUSION_4_CURATED_INPAINTING,
+    ]:
+        prompts = [
+            UcPrompt(
+                category="heavy",
+                name="heavy",
+                text="blurry, lowres, error, film grain, scan artifacts, worst quality, bad quality, jpeg artifacts, very displeasing, chromatic aberration, logo, dated, signature, multiple views, gigantic breasts"
+            ),
+            UcPrompt(
+                category="light",
+                name="light",
+                text="blurry, lowres, error, worst quality, bad quality, jpeg artifacts, very displeasing, logo, dated, signature"
+            ),
+            UcPrompt(
+                category="none",
+                name="none",
+                text=""
+            ),
+        ]
+    elif model in [
+        Model.CUSTOM,
+        Model.NAI_DIFFUSION_4_FULL,
+        Model.NAI_DIFFUSION_4_FULL_INPAINTING,
     ]:
         prompts = [
-            UcPrompt(category="heavy", name="heavy",
-                     text="blurry, lowres, error, film grain, scan artifacts, worst quality, bad quality, jpeg artifacts, very displeasing, chromatic aberration, logo, dated, signature, multiple views, gigantic breasts"),
-            UcPrompt(category="light", name="light",
-                     text="blurry, lowres, error, worst quality, bad quality, jpeg artifacts, very displeasing, logo, dated, signature"),
-            UcPrompt(category="none", name="none", text=""),
+            UcPrompt(
+                category="heavy",
+                name="heavy",
+                text="blurry, lowres, error, film grain, scan artifacts, worst quality, bad quality, jpeg artifacts, very displeasing, chromatic aberration, multiple views, logo, too many watermarks, white blank page, blank page"
+            ),
+            UcPrompt(
+                category="light",
+                name="light",
+                text="blurry, lowres, error, worst quality, bad quality, jpeg artifacts, very displeasing, white blank page, blank page"
+            ),
+            UcPrompt(
+                category="none",
+                name="none",
+                text=""
+            ),
         ]
     return prompts
 

{novelai_python-0.7.5 → novelai_python-0.7.6}/src/novelai_python/sdk/ai/generate_image/__init__.py

@@ -509,33 +509,13 @@ class GenerateImageInfer(ApiBaseModel):
         **NOTE: It is forbidden to write logic that does not belong to novelai.net here**
         :return: self
         """
-        if self.action == Action.GENERATE:
-            if self.parameters.image is not None:
-                raise ValueError("You are using generate action, image is not required for non-generate mode.")
-            if self.parameters.mask is not None:
-                raise ValueError("You are using generate action, mask is not required for non-generate mode.")
-
-        if self.action == Action.INFILL:
-            if self.model not in INPAINTING_MODEL_LIST:
-                raise ValueError(f"You must use {INPAINTING_MODEL_LIST}")
-            if not self.parameters.mask:
-                logger.warning("Mask maybe required for infill mode!")
-
-        if self.action == Action.IMG2IMG:
-            if self.parameters.extra_noise_seed is None:
-                self.parameters.extra_noise_seed = self.parameters.seed
-
-        if self.action == Action.IMG2IMG:
-            if self.parameters.sm_dyn is True:
-                logger.warning("sm_dyn be disabled when sm in Img2Img mode.")
-            if self.parameters.sm is True:
-                logger.warning("sm be disabled when sm_dyn in Img2Img mode.")
-            self.parameters.sm = False
-            self.parameters.sm_dyn = False
-            if self.parameters.image is None:
-                raise ValueError("image is must required for img2img mode.")
+        if self.parameters.characterPrompts:
+            self.parameters.use_coords = any([c.center != PositionMap.AUTO for c in self.parameters.characterPrompts])
+        # Mix Prompt Warning
+        if self.input.count('|') > 6:
+            logger.warning("Maximum prompt mixes exceeded. Extra will be ignored.")
 
-        # Check Image
+        # Remove the image if the model does not support it
         if self.parameters.image is None:
             self.parameters.strength = None
             self.parameters.noise = None
@@ -553,31 +533,10 @@ class GenerateImageInfer(ApiBaseModel):
             self.parameters.sm_dyn = True
 
         if self.parameters.sm_dyn and (not self.parameters.sm):
-            self.parameters.sm = True
-
-        # AUTO SMEA SETTING
-        """
-        def sm_size_cast(model: Model):
-            if model in [
-                Model.NAI_DIFFUSION_V3,
-                Model.NAI_DIFFUSION_V3_INPAINTING,
-                Model.NAI_DIFFUSION_FURRY_V3,
-                Model.NAI_DIFFUSION_FURRY_V3_INPAINTING,
-                Model.CUSTOM
-            ]:
-                return 2166785
-            return 1064961
-        if self.parameters.autoSmea and (self.parameters.width * self.parameters.height < sm_size_cast(self.model)):
-            self.parameters.autoSmea = False
-        """
-
-        if self.parameters.characterPrompts:
-            self.parameters.use_coords = any([c.center != PositionMap.AUTO for c in self.parameters.characterPrompts])
-
-        # Mix Prompt Warning
-        if self.input.count('|') > 6:
-            logger.warning("Maximum prompt mixes exceeded. Extra will be ignored.")
+            logger.warning("sm_dyn is disabled when sm is disabled.")
+            self.parameters.sm_dyn = False
 
+        # DDIM DON'T SUPPORT sm and sm_dyn
         if self.parameters.sampler in [
             Sampler.DDIM,
             Sampler.DDIM_V3
@@ -592,17 +551,33 @@ class GenerateImageInfer(ApiBaseModel):
 
         if self.parameters.uncond_scale == 0:
             self.parameters.uncond_scale = 0.00001
-
         if get_model_group(self.model) == ModelGroups.STABLE_DIFFUSION:
             self.parameters.uncond_scale = None
 
+        # Remove noise_schedule and cfg_rescale if not supported
         if not get_supported_params(self.model).noiseSchedule:
-            self.parameters.noise_schedule = None
             self.parameters.cfg_rescale = None
+            self.parameters.noise_schedule = None
 
         if (self.model in INPAINTING_MODEL_LIST) and (self.parameters.sampler in [Sampler.DDIM, Sampler.DDIM_V3]):
             self.parameters.sampler = Sampler.K_EULER_ANCESTRAL
 
+        # AUTO SMEA SETTING
+        """
+        def sm_size_cast(model: Model):
+            if model in [
+                Model.NAI_DIFFUSION_V3,
+                Model.NAI_DIFFUSION_V3_INPAINTING,
+                Model.NAI_DIFFUSION_FURRY_V3,
+                Model.NAI_DIFFUSION_FURRY_V3_INPAINTING,
+                Model.CUSTOM
+            ]:
+                return 2166785
+            return 1064961
+        if self.parameters.autoSmea and (self.parameters.width * self.parameters.height < sm_size_cast(self.model)):
+            self.parameters.autoSmea = False
+        """
+
         if self.parameters.sampler in [
             Sampler.DDIM,
             Sampler.PLMS,
@@ -638,16 +613,46 @@ class GenerateImageInfer(ApiBaseModel):
             self.parameters.sm = None
             self.parameters.sm_dyn = None
 
+        if not get_supported_params(self.model).dynamicThresholding:
+            self.parameters.dynamic_thresholding = False
+
+        if self.model == Model.CUSTOM and self.action == Action.INFILL:
+            self.endpoint = "http://custom-exposed-full-inpaint-anime-v4-staging.ai.svc.cluster.local/"
+
         # == Noise Schedule ==
         if self.parameters.noise_schedule is None:
             self.parameters.noise_schedule = get_default_noise_schedule(self.parameters.sampler)
-        supported_noise_schedule = get_supported_noise_schedule(self.parameters.sampler)
+        supported_noise_schedule = get_supported_noise_schedule(sample_type=self.parameters.sampler, model=self.model)
         if supported_noise_schedule:
             if self.parameters.noise_schedule not in supported_noise_schedule:
                 raise ValueError(f"Invalid noise_schedule, must be one of {supported_noise_schedule}")
         else:
             logger.warning(f"Inactive sampler {self.parameters.sampler} does not support noise_schedule.")
 
+        # == Action ==
+        if self.action == Action.GENERATE:
+            if self.parameters.image is not None:
+                raise ValueError("You are using generate action, image is not required for non-generate mode.")
+            if self.parameters.mask is not None:
+                raise ValueError("You are using generate action, mask is not required for non-generate mode.")
+        if self.action == Action.INFILL:
+            if self.model not in INPAINTING_MODEL_LIST:
+                raise ValueError(f"You must use {INPAINTING_MODEL_LIST}")
+            if not self.parameters.mask:
+                logger.warning("Mask maybe required for infill mode!")
+        if self.action == Action.IMG2IMG:
+            if self.parameters.extra_noise_seed is None:
+                self.parameters.extra_noise_seed = self.parameters.seed
+        if self.action == Action.IMG2IMG:
+            if self.parameters.sm_dyn is True:
+                logger.warning("sm_dyn be disabled when sm in Img2Img mode.")
+            if self.parameters.sm is True:
+                logger.warning("sm be disabled when sm_dyn in Img2Img mode.")
+            self.parameters.sm = False
+            self.parameters.sm_dyn = False
+            if self.parameters.image is None:
+                raise ValueError("image is must required for img2img mode.")
+
         return self
 
     @property

{novelai_python-0.7.5 → novelai_python-0.7.6}/src/novelai_python/tool/image_metadata/__init__.py

@@ -12,7 +12,7 @@ from nacl.encoding import Base64Encoder
 from nacl.signing import VerifyKey
 from pydantic import BaseModel, ConfigDict
 
-from novelai_python.sdk.ai._enum import PROMOTION, Model
+from novelai_python.sdk.ai._enum import PROMOTION, Model, find_model_by_hashcode
 from .bch_utils import fec_decode
 from .lsb_extractor import ImageLsbDataExtractor
 from .lsb_injector import inject_data
@@ -125,7 +125,7 @@ class ImageMetadata(BaseModel):
         Get the model used to generate the image from the Model code
         :return:  Model or None
         """
-        return PROMOTION.get(self.Source, None)
+        return find_model_by_hashcode(self.Source)
 
     @staticmethod
     def reset_alpha(image: Image.Image) -> BytesIO: