Merge pull request #642 from zanllp/feat/add-novelai-support

feat: Add support for NovelAI

Author: zanllp

README.md

@@ -20,7 +20,7 @@ https://github.com/zanllp/sd-webui-infinite-image-browsing/assets/25872019/807b8
 | Stable Diffusion web UI| Supported        | --                   |
 | ComfyUI                | Partially supported | --                |
 | Fooocus                | Supported        | --                   |
-| NovelAI                | Not supported    | Planned              |
+| NovelAI                | Supported    | --            |
 | StableSwarmUI          | Not supported    | Planned              |
 
 

javascript/index.js

@@ -13,7 +13,7 @@ Promise.resolve().then(async () => {
     <link rel="icon" href="/favicon.ico" />
     <meta name="viewport" content="width=device-width, initial-scale=1.0" />
     <title>Infinite Image Browsing</title>
-    <script type="module" crossorigin src="/infinite_image_browsing/fe-static/assets/index-dfe671d9.js"></script>
+    <script type="module" crossorigin src="/infinite_image_browsing/fe-static/assets/index-e8d50841.js"></script>
     <link rel="stylesheet" href="/infinite_image_browsing/fe-static/assets/index-f6de8b10.css">
   </head>
 

scripts/iib/parsers/index.py

@@ -1,14 +1,15 @@
 from scripts.iib.parsers.comfyui import ComfyUIParser
 from scripts.iib.parsers.sd_webui import SdWebUIParser
 from scripts.iib.parsers.fooocus import FooocusParser
+from scripts.iib.parsers.novelai import NovelAIParser
 from scripts.iib.parsers.model import ImageGenerationInfo
 from scripts.iib.logger import logger
 from PIL import Image
 import traceback
 
 
 def parse_image_info(image_path: str) -> ImageGenerationInfo:
-    parsers = [ComfyUIParser, FooocusParser, SdWebUIParser]
+    parsers = [ComfyUIParser, FooocusParser, NovelAIParser, SdWebUIParser]
     with Image.open(image_path) as img:
         for parser in parsers:
             if parser.test(img, image_path):

scripts/iib/parsers/novelai.py

@@ -0,0 +1,45 @@
+import json
+from PIL import Image
+
+from scripts.iib.tool import (
+    parse_generation_parameters,
+)
+from scripts.iib.parsers.model import ImageGenerationInfo, ImageGenerationParams
+
+def replace_punctuation(input_string):
+    return input_string.replace(',', ' ').replace('\n', ' ')
+
+class NovelAIParser:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def parse(clz, img, file_path):
+        info = ""
+        params = None
+        if not clz.test(img, file_path):
+            raise Exception("The input image does not match the current parser.")
+        data = json.loads(img.info.get('Comment'))
+        meta_kv = [f"""Steps: {data["steps"]}, Source Identifier: NovelAI"""]
+        for key, value in data.items():
+            if key not in ["prompt"]:
+                value = replace_punctuation(str(value))
+                meta_kv.append(f"{key}: {value}")
+        meta = ', '.join(meta_kv)
+        info = data["prompt"] + '\n' + meta
+
+        params = parse_generation_parameters(info)
+
+        return ImageGenerationInfo(
+            info,
+            ImageGenerationParams(
+                meta=params["meta"], pos_prompt=params["pos_prompt"]
+            ),
+        )
+
+    @classmethod
+    def test(clz, img: Image, file_path: str) -> bool:
+        try:
+            return img.info.get('Software') == 'NovelAI' and isinstance(img.info.get('Comment'), str)
+        except Exception:
+            return False