Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Dev #24

Merged
merged 5 commits into from
Feb 13, 2024
Merged

Dev #24

Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion playground/mask/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,11 @@
with open('ori.png', 'rb') as f:
ori_bytes = f.read()

return_bytes = create_mask_from_sketch(original_img_bytes=ori_bytes, sketch_img_bytes=sk_bytes, jagged_edges=True)
return_bytes = create_mask_from_sketch(original_img_bytes=ori_bytes,
sketch_img_bytes=sk_bytes,
jagged_edges=True,
min_block_size=15
)

with open('mask_export.png', 'wb') as f:
f.write(return_bytes)
38 changes: 38 additions & 0 deletions playground/suggest_tag.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
# -*- coding: utf-8 -*-
# @Time : 2024/2/13 下午8:25
# @Author : sudoskys
# @File : suggest_tag.py
# @Software: PyCharm
import asyncio
import os

from dotenv import load_dotenv
from loguru import logger
from pydantic import SecretStr

from novelai_python import APIError
from novelai_python import SuggestTags, SuggestTagsResp, JwtCredential

load_dotenv()

token = None
jwt = os.getenv("NOVELAI_JWT") or token


async def main():
globe_s = JwtCredential(jwt_token=SecretStr(jwt))
try:
_res = await SuggestTags().request(
session=globe_s
)
_res: SuggestTagsResp
print(f"Information: {_res}")
print(_res.model_dump())
except APIError as e:
logger.exception(e)
print(e.__dict__)
return


loop = asyncio.get_event_loop()
loop.run_until_complete(main())
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[project]
name = "novelai-python"
version = "0.3.1"
version = "0.3.2"
description = "Novelai Python Binding With Pydantic"
authors = [
{ name = "sudoskys", email = "[email protected]" },
Expand Down
4 changes: 4 additions & 0 deletions src/novelai_python/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
)
from .credential import JwtCredential, LoginCredential, ApiCredential
from .sdk import GenerateImageInfer, ImageGenerateResp
from .sdk import SuggestTags, SuggestTagsResp
from .sdk import Information, InformationResp
from .sdk import Login, LoginResp
from .sdk import Subscription, SubscriptionResp
Expand All @@ -31,6 +32,9 @@
"Login",
"LoginResp",

"SuggestTags",
"SuggestTagsResp",

"Information",
"InformationResp",

Expand Down
9 changes: 9 additions & 0 deletions src/novelai_python/_response/ai/generate_image.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,3 +22,12 @@ def query_params(self, key: str, default=None):
if not isinstance(self.meta.raw_request.get("parameters"), dict):
raise Exception("Resp parameters is not dict")
return self.meta.raw_request.get("parameters").get(key, default)


class SuggestTagsResp(RespBase):
class Tag(BaseModel):
tag: str
count: int
confidence: float

tags: List[Tag] = None
1 change: 1 addition & 0 deletions src/novelai_python/credential/ApiToken.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ async def get_session(self, timeout: int = 180, update_headers: dict = None):
"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": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36",
"Authorization": f"Bearer {self.api_token.get_secret_value()}",
"Content-Type": "application/json",
"Origin": "https://novelai.net",
Expand Down
2 changes: 1 addition & 1 deletion src/novelai_python/credential/JwtToken.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ async def get_session(self, timeout: int = 180, update_headers: dict = None):
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",
"User-Agent": "Mozilla/5.0 (X11; Linux x86_64; rv:123.0) Gecko/20100101 Firefox/123.0",
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36",
"Accept-Encoding": "gzip, deflate, br",
"Authorization": f"Bearer {self.jwt_token.get_secret_value()}",
"Content-Type": "application/json",
Expand Down
2 changes: 1 addition & 1 deletion src/novelai_python/credential/UserAuth.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ async def get_session(self, timeout: int = 180, update_headers: dict = None):
resp = await Login.build(user_name=self.username, password=self.password.get_secret_value()).request()
self._session = AsyncSession(timeout=timeout, headers={
"Accept": "*/*",
"User-Agent": "Mozilla/5.0 (X11; Linux x86_64; rv:123.0) Gecko/20100101 Firefox/123.0",
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36",
"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",
"Authorization": f"Bearer {resp.accessToken}",
Expand Down
3 changes: 2 additions & 1 deletion src/novelai_python/sdk/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,8 @@
# @Software: PyCharm

from .ai.generate_image import GenerateImageInfer, ImageGenerateResp # noqa 401
from .ai.generate_image.suggest_tags import SuggestTags, SuggestTagsResp # noqa 401
from .ai.upscale import Upscale, UpscaleResp # noqa 401
from .user.information import Information, InformationResp # noqa 401
from .user.login import Login, LoginResp # noqa 401
from .user.subscription import Subscription, SubscriptionResp # noqa 401
from .ai.upscale import Upscale, UpscaleResp # noqa 401
Original file line number Diff line number Diff line change
@@ -1,14 +1,13 @@
# -*- coding: utf-8 -*-
# @Time : 2024/1/26 上午11:21
# @Time : 2024/2/13 下午8:08
# @Author : sudoskys
# @File : generate_image.py
# @File : __init__.py.py
# @Software: PyCharm
import base64
import json
import math
import random
from copy import deepcopy
from enum import Enum, IntEnum
from io import BytesIO
from typing import Optional, Union
from urllib.parse import urlparse
Expand All @@ -21,84 +20,12 @@
from pydantic import BaseModel, ConfigDict, PrivateAttr, field_validator, model_validator, Field
from typing_extensions import override

from ..schema import ApiBaseModel
from ..._exceptions import APIError, AuthError, ConcurrentGenerationError, SessionHttpError
from ..._response.ai.generate_image import ImageGenerateResp
from ...credential import CredentialBase
from ...utils import try_jsonfy, NovelAiMetadata


class Sampler(Enum):
K_EULER = "k_euler"
K_EULER_ANCESTRAL = "k_euler_ancestral"
K_DPMPP_2S_ANCESTRAL = "k_dpmpp_2s_ancestral"
K_DPMPP_2M = "k_dpmpp_2m"
K_DPMPP_SDE = "k_dpmpp_sde"
DDIM_V3 = "ddim_v3"


class NoiseSchedule(Enum):
NATIVE = "native"
KARRAS = "karras"
EXPONENTIAL = "exponential"
POLYEXPONENTIAL = "polyexponential"


class UCPreset(IntEnum):
TYPE0 = 0
TYPE1 = 1
TYPE2 = 2
TYPE3 = 3


class Action(Enum):
GENERATE = "generate"
"""Generate Image"""
IMG2IMG = "img2img"
"""Image to Image"""
INFILL = "infill"
"""Inpainting"""


class Model(Enum):
NAI_DIFFUSION_3 = "nai-diffusion-3"
NAI_DIFFUSION_3_INPAINTING = "nai-diffusion-3-inpainting"

NAI_DIFFUSION = "nai-diffusion"
NAI_DIFFUSION_INPAINTING = "nai-diffusion-inpainting"

SAFE_DIFFUSION = "safe-diffusion"
SAFE_DIFFUSION_INPAINTING = "safe-diffusion-inpainting"

NAI_DIFFUSION_FURRY = "nai-diffusion-furry"
FURRY_DIFFUSION_INPAINTING = "furry-diffusion-inpainting"


class ControlNetModel(Enum):
HED = "hed"
"""边缘检测"""
MIDAS = "midas"
"""景深"""
FAKE_SCRIBBLE = "fake_scribble"
"""伪涂鸦"""
M_LSD = "mlsd"
"""(建筑)线条检测"""
LANDSCAPER = "uniformer"
"""风景生成"""


class Resolution(Enum):
RES_512_768 = (512, 768)
RES_768_512 = (768, 512)
RES_640_640 = (640, 640)
RES_832_1216 = (832, 1216)
RES_1216_832 = (1216, 832)
RES_1024_1024 = (1024, 1024)
RES_1024_1536 = (1024, 1536)
RES_1536_1024 = (1536, 1024)
RES_1472_1472 = (1472, 1472)
RES_1088_1920 = (1088, 1920)
RES_1920_1088 = (1920, 1088)
from ._enum import Model, Sampler, NoiseSchedule, ControlNetModel, Action, UCPreset
from ...schema import ApiBaseModel
from ...._exceptions import APIError, AuthError, ConcurrentGenerationError, SessionHttpError
from ...._response.ai.generate_image import ImageGenerateResp
from ....credential import CredentialBase
from ....utils import try_jsonfy, NovelAiMetadata


class GenerateImageInfer(ApiBaseModel):
Expand Down Expand Up @@ -395,7 +322,7 @@ async def necessary_headers(self, request_data) -> dict:
"Sec-Fetch-Site": "same-site",
"Pragma": "no-cache",
"Cache-Control": "no-cache",
"User-Agent": "Mozilla/5.0 (X11; Linux x86_64; rv:123.0) Gecko/20100101 Firefox/123.0"
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36",
}

async def request(self,
Expand Down
78 changes: 78 additions & 0 deletions src/novelai_python/sdk/ai/generate_image/_enum.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
# -*- coding: utf-8 -*-
# @Time : 2024/2/13 下午8:10
# @Author : sudoskys
# @File : _enum.py
# @Software: PyCharm
from enum import Enum, IntEnum


class Sampler(Enum):
K_EULER = "k_euler"
K_EULER_ANCESTRAL = "k_euler_ancestral"
K_DPMPP_2S_ANCESTRAL = "k_dpmpp_2s_ancestral"
K_DPMPP_2M = "k_dpmpp_2m"
K_DPMPP_SDE = "k_dpmpp_sde"
DDIM_V3 = "ddim_v3"


class NoiseSchedule(Enum):
NATIVE = "native"
KARRAS = "karras"
EXPONENTIAL = "exponential"
POLYEXPONENTIAL = "polyexponential"


class UCPreset(IntEnum):
TYPE0 = 0
TYPE1 = 1
TYPE2 = 2
TYPE3 = 3


class Action(Enum):
GENERATE = "generate"
"""Generate Image"""
IMG2IMG = "img2img"
"""Image to Image"""
INFILL = "infill"
"""Inpainting"""


class ControlNetModel(Enum):
HED = "hed"
"""边缘检测"""
MIDAS = "midas"
"""景深"""
FAKE_SCRIBBLE = "fake_scribble"
"""伪涂鸦"""
M_LSD = "mlsd"
"""(建筑)线条检测"""
LANDSCAPER = "uniformer"
"""风景生成"""


class Resolution(Enum):
RES_512_768 = (512, 768)
RES_768_512 = (768, 512)
RES_640_640 = (640, 640)
RES_832_1216 = (832, 1216)
RES_1216_832 = (1216, 832)
RES_1024_1024 = (1024, 1024)
RES_1024_1536 = (1024, 1536)
RES_1536_1024 = (1536, 1024)
RES_1472_1472 = (1472, 1472)
RES_1088_1920 = (1088, 1920)
RES_1920_1088 = (1920, 1088)

class Model(Enum):
NAI_DIFFUSION_3 = "nai-diffusion-3"
NAI_DIFFUSION_3_INPAINTING = "nai-diffusion-3-inpainting"

NAI_DIFFUSION = "nai-diffusion"
NAI_DIFFUSION_INPAINTING = "nai-diffusion-inpainting"

SAFE_DIFFUSION = "safe-diffusion"
SAFE_DIFFUSION_INPAINTING = "safe-diffusion-inpainting"

NAI_DIFFUSION_FURRY = "nai-diffusion-furry"
FURRY_DIFFUSION_INPAINTING = "furry-diffusion-inpainting"
Loading
Loading