Merge pull request #1 from LlmKira/dev

🔨 chore(generate_image): Clean up generate_image.py

Author: sudoskys

README.md

@@ -49,6 +49,22 @@ loop.run_until_complete(main())
 
 ```
 
+#### Random Prompt
+
+```python
+from novelai_python.utils.random_prompt import RandomPromptGenerator
+
+s = RandomPromptGenerator(nsfw_enabled=False).generate()
+print(s)
+```
+
+#### Run A Server
+
+```shell
+pip install novelai_python
+python3 -m novelai_python.server -h '0.0.0.0' -p 7888
+```
+
 ## Acknowledgements 🙏
 
 [BackEnd](https://api.novelai.net/docs)

pdm.lock

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.1.5"
+version = "0.1.6"
 description = "Novelai Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "[email protected]" },
@@ -14,6 +14,8 @@ dependencies = [
     "shortuuid>=1.0.11",
     "Pillow>=10.2.0",
     "curl-cffi>=0.5.10",
+    "fastapi>=0.109.0",
+    "uvicorn[standard]>=0.27.0.post1",
 ]
 requires-python = ">=3.8"
 readme = "README.md"

pyproject.toml

@@ -22,8 +22,11 @@ class APIError(NovelAiError):
     code: Optional[str] = None
     response: Union[Dict[str, Any], str] = None
 
-    def __init__(self, message: str, request: Any, response: Any, status_code: str) -> None:
-        super().__init__(message)
+    def __init__(self, message: str, request: dict, response: Union[dict, str], status_code: str) -> None:
+        if not isinstance(response, dict):
+            response = {"error": f"data type error, should be dict, but got {type(response)}"}
+        if not isinstance(request, dict):
+            request = {"error": f"data type error, should be dict, but got {type(request)}"}
         self.request = request
         self.message = message
         self.code = status_code
@@ -35,5 +38,5 @@ class AuthError(APIError):
     AuthError is raised when the API returns an error.
     """
 
-    def __init__(self, message: str, request: Any, response: Any, status_code: str) -> None:
+    def __init__(self, message: str, request: dict, response: Union[dict, str], status_code: str) -> None:
         super().__init__(message, request, response, status_code)

src/novelai_python/_exceptions.py

@@ -4,7 +4,8 @@
 # @File    : __init__.py.py
 # @Software: PyCharm
 from .JwtToken import JwtCredential
-
+from pydantic import SecretStr
 __all__ = [
-    "JwtCredential"
+    "JwtCredential",
+    "SecretStr"
 ]

src/novelai_python/credential/__init__.py

@@ -6,13 +6,13 @@
 import json
 import random
 from io import BytesIO
-from typing import Optional, Union
+from typing import Optional, Union, Literal
 from zipfile import ZipFile
 
 import httpx
 from curl_cffi.requests import AsyncSession
 from loguru import logger
-from pydantic import BaseModel, ConfigDict, PrivateAttr, field_validator
+from pydantic import BaseModel, ConfigDict, PrivateAttr, field_validator, model_validator
 
 from ..._exceptions import APIError, AuthError
 from ..._response import ImageGenerateResp
@@ -30,6 +30,7 @@ class Params(BaseModel):
         controlnet_strength: Optional[int] = 1
         dynamic_thresholding: Optional[bool] = False
         height: Optional[int] = 1216
+        image: Optional[str] = None  # img2img,base64
         legacy: Optional[bool] = False
         legacy_v3_extend: Optional[bool] = False
         n_samples: Optional[int] = 1
@@ -41,7 +42,7 @@ class Params(BaseModel):
             "extra digit, fewer digits, cropped,                        worst quality, low quality, normal quality, "
             "jpeg artifacts, signature, watermark, username, blurry"
         )
-        noise_schedule: Optional[str] = "native"
+        noise_schedule: Optional[Union[str, Literal['native', 'polyexponential', 'exponential']]] = "native"
         params_version: Optional[int] = 1
         qualityToggle: Optional[bool] = True
         sampler: Optional[str] = "k_euler"
@@ -60,15 +61,25 @@ def seed_validator(cls, v: int):
                 v = random.randint(0, 2 ** 32 - 1)
             return v
 
+        @model_validator(mode="after")
+        def validate_img2img(self):
+
+            image = True if self.image else False
+            add_origin = True if self.add_original_image else False
+            if image != add_origin:
+                raise ValueError('Invalid Model Params For img2img2 mode... image should match add_original_image!')
+            return self
+
         @field_validator('sampler')
         def sampler_validator(cls, v: str):
-            if v not in ["k_euler", "k_euler_ancestral", "k_dpmpp_2m", "k_dpmpp_sde", "ddim_v3"]:
+            if v not in ["k_euler", "k_euler_ancestral", 'k_dpmpp_2s_ancestral', "k_dpmpp_2m", "k_dpmpp_sde",
+                         "ddim_v3"]:
                 raise ValueError("Invalid sampler.")
             return v
 
         @field_validator('noise_schedule')
         def noise_schedule_validator(cls, v: str):
-            if v not in ["native", "exponential", "polyexponential"]:
+            if v not in ["native", "karras", "exponential", "polyexponential"]:
                 raise ValueError("Invalid noise_schedule.")
             return v
 
@@ -111,7 +122,7 @@ def n_samples_validator(cls, v: int):
                 raise ValueError("Invalid n_samples, must be less than 8.")
             return v
 
-    action: Optional[str] = "generate"
+    action: Union[str, Literal["generate", "img2img"]] = "generate"
     input: str = "1girl, best quality, amazing quality, very aesthetic, absurdres"
     model: Optional[str] = "nai-diffusion-3"
     parameters: Params = Params()
@@ -162,6 +173,8 @@ def validate_charge(self):
     def build(cls,
               prompt: str,
               *,
+              model: str = "nai-diffusion-3",
+              action: Literal['generate', 'img2img'] = 'generate',
               negative_prompt: Optional[str] = None,
               override_negative_prompt: bool = False,
               seed: int = -1,
@@ -175,6 +188,8 @@ def build(cls,
         正负面, step, cfg, 采样方式, seed
         :param override_negative_prompt:
         :param prompt:
+        :param model:
+        :param action: Mode for img generate
         :param negative_prompt:
         :param seed:
         :param steps:
@@ -209,13 +224,15 @@ def build(cls,
         param = {k: v for k, v in param.items() if v is not None}
         return cls(
             input=prompt,
+            model=model,
+            action=action,
             parameters=cls.Params(**param)
         )
 
     async def generate(self, session: Union[AsyncSession, JwtCredential]) -> ImageGenerateResp:
         if isinstance(session, JwtCredential):
             session = session.session
-        request_data = self.model_dump()
+        request_data = self.model_dump(exclude_none=True)
         logger.debug(f"Request Data: {request_data}")
         try:
             assert hasattr(session, "post"), "session must have post method."

src/novelai_python/sdk/ai/generate_image.py

@@ -0,0 +1,82 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/1/30 下午11:05
+# @Author  : sudoskys
+# @File    : server.py
+# @Software: PyCharm
+import io
+import sys
+import zipfile
+
+import uvicorn
+from fastapi import FastAPI, Depends, Security
+from fastapi.security import APIKeyHeader
+from starlette.responses import JSONResponse, StreamingResponse
+
+from .credential import JwtCredential, SecretStr
+from .sdk.ai.generate_image import GenerateImageInfer
+
+app = FastAPI()
+token_key = APIKeyHeader(name="Authorization")
+session = {}
+
+
+def get_session(token: str):
+    if token not in session:
+        session[token] = JwtCredential(jwt_token=SecretStr(token))
+    return session[token]
+
+
+def get_current_token(auth_key: str = Security(token_key)):
+    return auth_key
+
+
+@app.get("/health")
+async def health():
+    return {"status": "ok"}
+
+
+@app.post("/ai/generate_image")
+async def generate_image(
+        req: GenerateImageInfer,
+        current_token: str = Depends(get_current_token)
+):
+    """
+    生成图片
+    :param current_token: Authorization
+    :param req: GenerateImageInfer
+    :return:
+    """
+    try:
+        _result = await req.generate(session=get_session(current_token))
+        zip_file_bytes = io.BytesIO()
+        with zipfile.ZipFile(zip_file_bytes, mode="w", compression=zipfile.ZIP_DEFLATED) as zip_file:
+            for file in _result.files:
+                zip_file.writestr(zinfo_or_arcname=file[0], data=file[1])
+        # return the zip file
+        zip_file_bytes.seek(0)
+        return StreamingResponse(zip_file_bytes, media_type='application/zip', headers={
+            'Content-Disposition': 'attachment;filename=multiple_files.zip'
+        })
+    except Exception as e:
+        return JSONResponse(status_code=500, content=e.__dict__)
+
+
+# 获取输入参数
+def usage():
+    print("Usage: python -m novelai_python.server -h <host> -p <port>")
+    sys.exit(0)
+
+
+if __name__ == '__main__':
+    import getopt
+
+    opts = {}
+    try:
+        opts, args = getopt.getopt(sys.argv[1:], "h:p:", ["host=", "port="])
+    except getopt.GetoptError:
+        usage()
+    opts = dict(opts)
+    server_host = opts.get("-h", "0.0.0.0")
+    server_port = int(opts.get("-p", 10087))
+    print(f"Docs: http://{server_host}:{server_port}/docs")
+    uvicorn.run(app, host=server_host, port=server_port)

src/novelai_python/server.py

@@ -7,6 +7,8 @@
 import json
 from typing import Union
 
+from loguru import logger
+
 from .hash import NovelAiMetadata
 
 
@@ -28,5 +30,6 @@ def try_jsonfy(obj: Union[str, dict, list, tuple]):
     """
     try:
         return json.loads(obj)
-    except json.JSONDecoder:
+    except Exception:
+        logger.error(f"Decode Error {obj}")
         return f"Decode Error {type(obj)}"

src/novelai_python/utils/__init__.py

@@ -0,0 +1,70 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/1/27 上午10:41
+# @Author  : sudoskys
+# @File    : test_random_prompt.py
+# @Software: PyCharm
+
+from src.novelai_python.utils.random_prompt import RandomPromptGenerator
+
+
+def test_generate_returns_string():
+    generator = RandomPromptGenerator(nsfw_enabled=True)
+    result = generator.generate()
+    assert isinstance(result, str)
+
+
+def test_generate_returns_non_empty_string():
+    generator = RandomPromptGenerator(nsfw_enabled=True)
+    result = generator.generate()
+    assert len(result) > 0
+
+
+def test_generate_returns_different_results():
+    generator = RandomPromptGenerator(nsfw_enabled=True)
+    result1 = generator.generate()
+    result2 = generator.generate()
+    assert result1 != result2
+
+
+def test_generate_with_nsfw_disabled():
+    generator = RandomPromptGenerator(nsfw_enabled=False)
+    result = generator.generate()
+    assert 'nsfw' not in result
+
+
+def test_generate_with_nsfw_enabled():
+    generator = RandomPromptGenerator(nsfw_enabled=True)
+    result = generator.generate()
+    assert 'nsfw' in result
+
+
+def test_get_weighted_choice_returns_string():
+    generator = RandomPromptGenerator(nsfw_enabled=True)
+    result = generator.get_weighted_choice([['tag1', 1], ['tag2', 2]], [])
+    assert isinstance(result, str)
+
+
+def test_get_weighted_choice_returns_valid_tag():
+    generator = RandomPromptGenerator(nsfw_enabled=True)
+    result = generator.get_weighted_choice([['tag1', 1], ['tag2', 2]], [])
+    assert result in ['tag1', 'tag2']
+
+
+def test_character_features_returns_list():
+    generator = RandomPromptGenerator(nsfw_enabled=True)
+    result = generator.character_features('m', 'front', True, 1)
+    assert isinstance(result, list)
+
+
+def test_character_features_returns_non_empty_list():
+    generator = RandomPromptGenerator(nsfw_enabled=True)
+    result = generator.character_features('m', 'front', True, 1)
+    assert len(result) > 0
+
+
+def test_character_features_with_different_genders():
+    generator = RandomPromptGenerator(nsfw_enabled=True)
+    result_m = generator.character_features('m', 'front', True, 1)
+    result_f = generator.character_features('f', 'front', True, 1)
+    result_o = generator.character_features('o', 'front', True, 1)
+    assert result_m != result_f != result_o

tests/test_random_prompt.py

@@ -0,0 +1,28 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/1/30 下午11:52
+# @Author  : sudoskys
+# @File    : test_server_run.py
+# @Software: PyCharm
+from fastapi.testclient import TestClient
+
+from src.novelai_python.sdk.ai.generate_image import GenerateImageInfer
+from src.novelai_python.server import app, get_session
+
+client = TestClient(app)
+
+
+def test_health_check():
+    response = client.get("/health")
+    assert response.status_code == 200
+    assert response.json() == {"status": "ok"}
+
+
+def test_generate_image_with_valid_token():
+    valid_token = "valid_token"
+    get_session(valid_token)  # to simulate a valid session
+    response = client.post(
+        "/ai/generate_image",
+        headers={"Authorization": valid_token},
+        json=GenerateImageInfer(input="1girl").model_dump()
+    )
+    assert response.status_code == 500