Merge pull request #67 from LlmKira/dev

feat: Director Tools / imagetools

.nerve.toml

@@ -1,5 +1,5 @@
 # https://github.com/LlmKira/contributor/blob/main/.nerve.toml
-contributor = "fae6f549-fcd5-4c57-89ab-04c97ebb11cb"
+contributor = "9f30f440-3d09-43dd-aa61-285c66f89356"
 language = "English"
 issue_auto_label = true
 issue_title_format = true

playground/augment-image.py

@@ -0,0 +1,59 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/1/26 下午12:23
+# @Author  : sudoskys
+# @File    : __init__.py.py
+# @Software: PyCharm
+import asyncio
+import os
+import pathlib
+
+from dotenv import load_dotenv
+from pydantic import SecretStr
+
+from novelai_python import APIError, LoginCredential, JwtCredential, ImageGenerateResp
+from novelai_python import AugmentImageInfer
+from novelai_python.sdk.ai.augment_image import ReqType
+
+
+async def generate(
+        image,
+        request_type: ReqType = ReqType.SKETCH,
+):
+    jwt = os.getenv("NOVELAI_JWT", None)
+    if jwt is None:
+        raise ValueError("NOVELAI_JWT is not set in `.env` file, please create one and set it")
+    credential = JwtCredential(jwt_token=SecretStr(jwt))
+    """Or you can use the login credential to get the renewable jwt token"""
+    _login_credential = LoginCredential(
+        username=os.getenv("NOVELAI_USER"),
+        password=SecretStr(os.getenv("NOVELAI_PASS"))
+    )
+    try:
+        agent = AugmentImageInfer.build(
+            req_type=request_type,
+            image=image
+        )
+        # print(f"charge: {agent.calculate_cost(is_opus=True)} if you are vip3")
+        # print(f"charge: {agent.calculate_cost(is_opus=False)} if you are not vip3")
+        result = await agent.request(
+            session=credential
+        )
+    except APIError as e:
+        print(f"Error: {e.message}")
+        return None
+    else:
+        print(f"Meta: {result.meta}")
+    _res: ImageGenerateResp
+    file = result.files[0]
+    with open(f"{pathlib.Path(__file__).stem}.png", "wb") as f:
+        f.write(file[1])
+
+
+load_dotenv()
+loop = asyncio.new_event_loop()
+loop.run_until_complete(
+    generate(
+        image=pathlib.Path(__file__).parent / "static_image.png",
+        request_type=ReqType.SKETCH
+    )
+)

playground/generate_image.py

@@ -53,5 +53,5 @@ async def generate(prompt="1girl, year 2023, dynamic angle, best quality, amazin
 
 
 load_dotenv()
-loop = asyncio.get_event_loop()
+loop = asyncio.new_event_loop()
 loop.run_until_complete(generate())

playground/generate_stream.py

@@ -5,7 +5,7 @@
 from loguru import logger
 from pydantic import SecretStr
 
-from novelai_python import APIError, Login, LoginCredential
+from novelai_python import APIError, LoginCredential
 from novelai_python import JwtCredential
 from novelai_python.sdk.ai.generate_stream import TextLLMModel, LLMStream, LLMStreamResp
 
@@ -52,5 +52,5 @@ async def stream(prompt="Hello"):
         print(f"Meta: {loop_connect(_data)}")
 
 
-loop = asyncio.get_event_loop()
+loop = asyncio.new_event_loop()
 loop.run_until_complete(stream())

playground/generate_voice.py

@@ -1,6 +1,7 @@
 import asyncio
 import os
 import pathlib
+
 from dotenv import load_dotenv
 from pydantic import SecretStr
 
@@ -39,7 +40,7 @@ async def generate_voice(text: str):
         f.write(file)
 
 
-loop = asyncio.get_event_loop()
+loop = asyncio.new_event_loop()
 loop.run_until_complete(
     generate_voice("Hello, I am a test voice, limit 1000 characters")
 )

playground/other/augment-image.py

@@ -0,0 +1,61 @@
+import json
+
+import requests
+
+url = "https://image.novelai.net/ai/augment-image"
+
+payload = json.dumps({
+    "req_type": "bg-removal",
+    "width": 1024,
+    "height": 1024,
+    "image": "iVBORw0"
+})
+payload = json.dumps({
+    "req_type": "colorize",
+    "prompt": "",
+    "defry": 0,
+    "width": 2252,
+    "height": 465,
+    "image": "iVBORw0KGgoAAA"
+})
+payload = json.dumps({
+    "req_type": "lineart",
+    "width": 2252,
+    "height": 465,
+    "image": "iVBO"
+})
+payload = json.dumps({
+    "req_type": "sketch",
+    "width": 2252,
+    "height": 465,
+    "image": "iVBOR"
+})
+
+payload = json.dumps({
+    "req_type": "emotion",
+    "prompt": "neutral;;123132",
+    "defry": 0,
+    "width": 2252,
+    "height": 465,
+    "image": "iVBORw0KG"
+})
+payload = json.dumps({
+    "req_type": "declutter",
+    "width": 2252,
+    "height": 465,
+    "image": "iVBORw0"
+})
+APIKEY = "YOUR_API"
+
+headers = {
+    'User-Agent': "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36",
+    'Accept-Encoding': "gzip, deflate, br, zstd",
+    'Content-Type': "application/json",
+    'authorization': f"Bearer {APIKEY}",
+    'origin': "https://novelai.net",
+    'priority': "u=1, i"
+}
+
+response = requests.post(url, data=payload, headers=headers)
+
+print(response.text)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.4.11"
+version = "0.4.12"
 description = "NovelAI Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "[email protected]" },

record/ai/augment_image/emotion/schema.json

@@ -0,0 +1,8 @@
+{
+  "req_type": "emotion",
+  "prompt": "neutral;;",
+  "defry": 0,
+  "width": 1024,
+  "height": 1024,
+  "image": "Base64 Data"
+}

record/ai/augment_image/export.py

@@ -0,0 +1,88 @@
+import base64
+import json
+import pathlib
+from io import BytesIO
+
+from PIL import Image
+from loguru import logger
+
+
+def ignore(*args, **kwargs):
+    pass
+
+
+def decode_base64_in_dict(data, current_path):
+    if isinstance(data, dict):
+        for k, v in data.items():
+            if isinstance(v, dict) or isinstance(v, list):
+                decode_base64_in_dict(v, current_path)
+            if isinstance(v, list):
+                _new_list = []
+                for index, item in enumerate(v):
+                    if isinstance(item, str) and len(item) > 100:
+                        try:
+                            # Base64解码
+                            image_bytes = base64.b64decode(item)
+                            image = Image.open(BytesIO(image_bytes))
+                        except Exception as e:
+                            ignore(e)
+                        else:
+                            logger.info(f"Decoding Base64 data in {k}")
+                            img_name = f"{current_path}/{index}-{k}.png"
+                            image.save(img_name)
+                        _new_list.append('Base64 Data')
+                if _new_list:
+                    data[k] = _new_list
+            if isinstance(v, str) and len(v) > 100:
+                try:
+                    # Base64解码
+                    image_bytes = base64.b64decode(v)
+                    image = Image.open(BytesIO(image_bytes))
+                except Exception as e:
+                    ignore(e)
+                else:
+                    logger.info(f"Decoding Base64 data in {k}")
+                    img_name = f"{current_path}/{k}.png"
+                    image.save(img_name)
+                    data[k] = 'Base64 Data'
+    elif isinstance(data, list):
+        for item in data:
+            if isinstance(item, dict) or isinstance(item, list):
+                decode_base64_in_dict(item, current_path)
+    return data
+
+
+def handle_file(filename):
+    filename_wo_ext = filename.stem
+    pathlib.Path(filename_wo_ext).mkdir(parents=True, exist_ok=True)
+    with open(filename, 'r') as file:
+        json_data = json.load(file)
+    # 取消 headers 里面 Authorization 字段，然后写回
+    if 'headers' in json_data:
+        if 'Authorization' in json_data['headers']:
+            json_data['headers']['Authorization'] = 'Secret'
+            # 写回原文件
+            with open(filename, 'w') as file:
+                json.dump(json_data, file, indent=2)
+        if 'authorization' in json_data['headers']:
+            json_data['headers']['authorization'] = 'Secret'
+            # 写回原文件
+            with open(filename, 'w') as file:
+                json.dump(json_data, file, indent=2)
+    request_data = json.loads(json_data.get("body", ""))
+    request_data = decode_base64_in_dict(request_data, filename_wo_ext)
+    # 写出包含替换字段的 JSON 文件回同名的文件夹
+    with open(f"{filename_wo_ext}/schema.json", 'w') as jsonfile:
+        json.dump(request_data, jsonfile, indent=2)
+
+
+def main():
+    # 列出当前文件夹内所有的 .json 文件
+    json_files = pathlib.Path('.').glob('*.json')
+    for file in json_files:
+        logger.info(f"Handling {file}")
+        handle_file(file)
+
+
+if __name__ == "__main__":
+    main()

record/ai/generate_image/enhance.json

@@ -3,7 +3,6 @@
   "headers": {
     "User-Agent": "Mozilla/5.0 (X11; Linux x86_64; rv:123.0) Gecko/20100101 Firefox/123.0",
     "Accept": "*/*",
-    "Accept-Language": "zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2",
     "Content-Type": "application/json",
     "Authorization": "Secret",
     "Sec-Fetch-Dest": "empty",

record/ai/generate_image/image2image.json

@@ -3,7 +3,7 @@
   "headers": {
     "User-Agent": "Mozilla/5.0 (X11; Linux x86_64; rv:123.0) Gecko/20100101 Firefox/123.0",
     "Accept": "*/*",
-    "Accept-Language": "zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2",
+
     "Content-Type": "application/json",
     "Authorization": "Secret",
     "Sec-Fetch-Dest": "empty",

record/ai/generate_image/inpaint.json

@@ -3,7 +3,7 @@
   "headers": {
     "User-Agent": "Mozilla/5.0 (X11; Linux x86_64; rv:123.0) Gecko/20100101 Firefox/123.0",
     "Accept": "*/*",
-    "Accept-Language": "zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2",
+
     "Content-Type": "application/json",
     "Authorization": "Secret",
     "Sec-Fetch-Dest": "empty",

record/ai/generate_image/text2image.json

@@ -3,7 +3,7 @@
   "headers": {
     "User-Agent": "Mozilla/5.0 (X11; Linux x86_64; rv:123.0) Gecko/20100101 Firefox/123.0",
     "Accept": "*/*",
-    "Accept-Language": "zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2",
+
     "Content-Type": "application/json",
     "Authorization": "Secret",
     "Sec-Fetch-Dest": "empty",

record/ai/generate_image/vibe_img2img.json

@@ -3,7 +3,7 @@
     "headers": {
         "User-Agent": "Mozilla/5.0 (X11; Linux x86_64; rv:123.0) Gecko/20100101 Firefox/123.0",
         "Accept": "*/*",
-        "Accept-Language": "zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2",
+
         "Content-Type": "application/json",
         "Sec-Fetch-Dest": "empty",
         "Sec-Fetch-Mode": "cors",

record/ai/generate_image/vibe_inpaint.json

@@ -3,7 +3,7 @@
     "headers": {
         "User-Agent": "Mozilla/5.0 (X11; Linux x86_64; rv:123.0) Gecko/20100101 Firefox/123.0",
         "Accept": "*/*",
-        "Accept-Language": "zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2",
+
         "Content-Type": "application/json",
         "Sec-Fetch-Dest": "empty",
         "Sec-Fetch-Mode": "cors",

src/novelai_python/__init__.py

@@ -11,6 +11,7 @@
     SessionHttpError
 )
 from .credential import JwtCredential, LoginCredential, ApiCredential
+from .sdk import AugmentImageInfer
 from .sdk import GenerateImageInfer, ImageGenerateResp
 from .sdk import Information, InformationResp
 from .sdk import LLM, LLMResp
@@ -31,6 +32,8 @@
     "GenerateImageInfer",
     "ImageGenerateResp",
 
+    "AugmentImageInfer",
+
     "VoiceGenerate",
     "VoiceResponse",
 

src/novelai_python/credential/ApiToken.py

@@ -3,6 +3,7 @@
 # @Author  : sudoskys
 # @File    : ApiToken.py
 # @Software: PyCharm
+
 from curl_cffi.requests import AsyncSession
 from loguru import logger
 from pydantic import SecretStr, Field, field_validator
@@ -12,27 +13,26 @@
 
 class ApiCredential(CredentialBase):
     """
-    JwtCredential is the base class for all credential.
+    ApiCredential is the base class for all credential.
     """
     api_token: SecretStr = Field(None, description="api token")
-    _session: AsyncSession = None
 
     async def get_session(self, timeout: int = 180, update_headers: dict = None):
-        if update_headers is None:
-            update_headers = {}
-        if not self._session:
-            self._session = AsyncSession(timeout=timeout, headers={
-                "Accept": "*/*",
-                "Accept-Language": "zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2",
-                "Accept-Encoding": "gzip, deflate, br",
-                "User-Agent": FAKE_UA.edge,
-                "Authorization": f"Bearer {self.api_token.get_secret_value()}",
-                "Content-Type": "application/json",
-                "Origin": "https://novelai.net",
-                "Referer": "https://novelai.net/",
-            }, impersonate="edge101")
-        self._session.headers.update(update_headers)
-        return self._session
+        headers = {
+            "Accept": "*/*",
+            "Accept-Encoding": "gzip, deflate, br",
+            "User-Agent": FAKE_UA.edge,
+            "Authorization": f"Bearer {self.api_token.get_secret_value()}",
+            "Content-Type": "application/json",
+            "Origin": "https://novelai.net",
+            "Referer": "https://novelai.net/",
+        }
+
+        if update_headers:
+            assert isinstance(update_headers, dict), "update_headers must be a dict"
+            headers.update(update_headers)
+
+        return AsyncSession(timeout=timeout, headers=headers, impersonate="edge101")
 
     @field_validator('api_token')
     def check_api_token(cls, v: SecretStr):