Merge pull request #31 from LlmKira/dev

Update SDK

.gitignore

@@ -161,3 +161,17 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
 /playground/generate_image.png
+/playground/generate_image_img2img.png
+/playground/vibe_inpaint.png
+/playground/vibe_img2img.png
+/playground/vibe.png
+/playground/upscale.py.png
+/playground/upscale.png
+/playground/test.png
+/playground/random_play/
+/playground/upscale/
+/playground/vibe/
+/playground/mask/
+/playground/newtag/
+/playground/oldtag/
+/playground/art_assert/

playground/generate_image.py

@@ -5,7 +5,7 @@
 # @Software: PyCharm
 import asyncio
 import os
-
+import pathlib
 from dotenv import load_dotenv
 from pydantic import SecretStr
 
@@ -14,41 +14,42 @@
 from novelai_python.sdk.ai.generate_image import Action, Sampler
 from novelai_python.utils.useful import enum_to_list
 
-load_dotenv()
-
-enhance = "year 2023,dynamic angle, best quality, amazing quality, very aesthetic, absurdres"
-token = None
-jwt = os.getenv("NOVELAI_JWT") or token
 
-
-async def main():
-    globe_s = JwtCredential(jwt_token=SecretStr(jwt))
-    _res = await Login.build(user_name=os.getenv("NOVELAI_USER"), password=os.getenv("NOVELAI_PASS")
-                             ).request()
+async def generate(prompt="1girl, year 2023, dynamic angle, best quality, amazing quality, very aesthetic, absurdres"):
+    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 = Login.build(
+        user_name=os.getenv("NOVELAI_USER"),
+        password=os.getenv("NOVELAI_PASS")
+    )
+    # await _login_credential.request()
+    print(f"Action List:{enum_to_list(Action)}")
     try:
-        print(f"Action List:{enum_to_list(Action)}")
-        gen = GenerateImageInfer.build(
-            prompt=f"1girl, winter, jacket, sfw, angel, flower,{enhance}",
+        agent = GenerateImageInfer.build(
+            prompt=prompt,
             action=Action.GENERATE,
-            sampler=Sampler.DDIM_V3
+            sampler=Sampler.K_DPMPP_SDE,
+            qualityToggle=True,
         )
-        cost = gen.calculate_cost(is_opus=True)
-        print(f"charge: {cost} if you are vip3")
-        print(f"charge: {gen.calculate_cost(is_opus=False)} if you are not vip3")
-        _res = await gen.request(
-            session=globe_s, remove_sign=True
+        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, remove_sign=True
         )
     except APIError as e:
-        print(str(e))
-        print(e.response)
-        return
-
+        print(f"Error: {e.message}")
+        return None
+    else:
+        print(f"Meta: {result.meta}")
     _res: ImageGenerateResp
-    print(_res.meta)
-    file = _res.files[0]
-    with open("generate_image.png", "wb") as f:
+    file = result.files[0]
+    with open(f"{pathlib.Path(__file__).stem}.png", "wb") as f:
         f.write(file[1])
 
 
+load_dotenv()
 loop = asyncio.get_event_loop()
-loop.run_until_complete(main())
+loop.run_until_complete(generate())

playground/generate_image_img2img.py

@@ -4,57 +4,64 @@
 # @File    : generate_image_img2img.py
 # @Software: PyCharm
 import asyncio
-import base64
 import os
-
+import pathlib
 from dotenv import load_dotenv
 from pydantic import SecretStr
 
 from novelai_python import APIError, Login
 from novelai_python import GenerateImageInfer, ImageGenerateResp, JwtCredential
-from novelai_python.sdk.ai.generate_image import Action
-
-load_dotenv()
+from novelai_python.sdk.ai.generate_image import Action, Sampler
+from novelai_python.utils.useful import enum_to_list
 
-enhance = "year 2023,dynamic angle, best quality, amazing quality, very aesthetic, absurdres"
-token = None
-jwt = os.getenv("NOVELAI_JWT") or token
 
-
-async def main():
-    globe_s = JwtCredential(jwt_token=SecretStr(jwt))
-    _res = await Login.build(user_name=os.getenv("NOVELAI_USER"), password=os.getenv("NOVELAI_PASS")
-                             ).request()
-    with open("raw_test_image.png", "rb") as f:
-        data = f.read()
-    # Base64 encode the data
-    encoded = base64.b64encode(data).decode()
+async def generate(
+        prompt="1girl, year 2023, dynamic angle, best quality, amazing quality, very aesthetic, absurdres",
+        image_path="static_image.png"
+):
+    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 jwt token"""
+    _login_credential = Login.build(
+        user_name=os.getenv("NOVELAI_USER"),
+        password=os.getenv("NOVELAI_PASS")
+    )
+    # await _login_credential.request()
+    print(f"Action List:{enum_to_list(Action)}")
+    print(f"Image Path: {image_path}")
     try:
-        gen = GenerateImageInfer.build(
-            prompt=f"1girl, spring, jacket, sfw, angel, flower,{enhance}",
+        if not os.path.exists(image_path):
+            raise ValueError(f"Image not found: {image_path}")
+        with open(image_path, "rb") as f:
+            image = f.read()
+        agent = GenerateImageInfer.build(
+            prompt=prompt,
             action=Action.IMG2IMG,
-            image=encoded,
-            add_original_image=True,
-            strength=0.5,
-            width=1088,
-            height=896
+            sampler=Sampler.K_DPMPP_SDE,
+            image=image,
+            add_original_image=False,
+            strength=0.9,
+            noise=0.1,
+            qualityToggle=True,
         )
-        cost = gen.calculate_cost(is_opus=True)
-        print(f"charge: {cost} if you are vip3")
-        print(f"charge: {gen.calculate_cost(is_opus=True)}")
-        _res = await gen.request(
-            session=globe_s, remove_sign=True
+        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, remove_sign=True
         )
     except APIError as e:
-        print(e.response)
-        return
-
+        print(f"Error: {e.message}")
+        return None
+    else:
+        print(f"Meta: {result.meta}")
     _res: ImageGenerateResp
-    print(_res.meta)
-    file = _res.files[0]
-    with open("generate_image_img2img.png", "wb") as f:
+    file = result.files[0]
+    with open(f"{pathlib.Path(__file__).stem}.png", "wb") as f:
         f.write(file[1])
 
 
+load_dotenv()
 loop = asyncio.get_event_loop()
-loop.run_until_complete(main())
+loop.run_until_complete(generate(image_path="static_refer.png"))

playground/upscale.py

@@ -11,6 +11,7 @@
 # To run the demo, you need an event loop, for instance by using asyncio
 import asyncio
 import os
+from pathlib import Path
 
 from dotenv import load_dotenv
 from pydantic import SecretStr
@@ -26,11 +27,13 @@
 jwt = os.getenv("NOVELAI_JWT") or token
 
 
-async def main():
+async def generate(
+        image_path="static_refer.png"
+):
     globe_s = JwtCredential(jwt_token=SecretStr(jwt))
-    if not os.path.exists("generate_image.png"):
-        raise FileNotFoundError("generate_image.png not found")
-    with open("generate_image.png", "rb") as f:
+    if not os.path.exists(image_path):
+        raise FileNotFoundError(f"{image_path} not found")
+    with open(image_path, "rb") as f:
         data = f.read()
     try:
         print(f"Action List:{enum_to_list(Action)}")
@@ -47,9 +50,9 @@ async def main():
     _res: UpscaleResp
     print(_res.meta.endpoint)
     file = _res.files
-    with open("upscale.py.png", "wb") as f:
+    with open(f"{Path(__file__).stem}.png", "wb") as f:
         f.write(file[1])
 
 
 loop = asyncio.get_event_loop()
-loop.run_until_complete(main())
+loop.run_until_complete(generate())

playground/vibe.py

@@ -0,0 +1,62 @@
+# -*- coding: utf-8 -*-
+import asyncio
+import os
+import pathlib
+from dotenv import load_dotenv
+from pydantic import SecretStr
+
+from novelai_python import APIError, Login
+from novelai_python import GenerateImageInfer, ImageGenerateResp, JwtCredential
+from novelai_python.sdk.ai.generate_image import Action, Sampler
+from novelai_python.utils.useful import enum_to_list
+
+
+async def generate(
+        prompt="1girl, year 2023, dynamic angle, best quality, amazing quality, very aesthetic, absurdres",
+        image_path="static_refer.png"
+):
+    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 jwt token"""
+    _login_credential = Login.build(
+        user_name=os.getenv("NOVELAI_USER"),
+        password=os.getenv("NOVELAI_PASS")
+    )
+    # await _login_credential.request()
+    print(f"Action List:{enum_to_list(Action)}")
+    try:
+        if not os.path.exists(image_path):
+            raise ValueError(f"Image not found: {image_path}")
+        with open(image_path, "rb") as f:
+            image = f.read()
+        agent = GenerateImageInfer.build(
+            prompt=prompt,
+            action=Action.GENERATE,
+            sampler=Sampler.K_DPMPP_SDE,
+            reference_image=image,
+            reference_strength=0.6,
+            reference_information_extracted=1,
+            add_original_image=True,  # This Not affect the vibe generation
+            qualityToggle=True,
+        )
+        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, remove_sign=True
+        )
+    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.get_event_loop()
+loop.run_until_complete(generate())

playground/vibe_img2img.py

@@ -0,0 +1,72 @@
+# -*- coding: utf-8 -*-
+import asyncio
+import os
+import pathlib
+from dotenv import load_dotenv
+from pydantic import SecretStr
+
+from novelai_python import APIError, Login
+from novelai_python import GenerateImageInfer, ImageGenerateResp, JwtCredential
+from novelai_python.sdk.ai.generate_image import Action, Sampler
+from novelai_python.utils.useful import enum_to_list
+
+
+async def generate(
+        prompt="1girl, year 2023, dynamic angle, best quality, amazing quality, very aesthetic, absurdres",
+        image_path="static_image.png",
+        reference_image_path="static_refer.png"
+):
+    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 jwt token"""
+    _login_credential = Login.build(
+        user_name=os.getenv("NOVELAI_USER"),
+        password=os.getenv("NOVELAI_PASS")
+    )
+    # await _login_credential.request()
+    print(f"Action List:{enum_to_list(Action)}")
+    try:
+        if not os.path.exists(image_path):
+            raise ValueError(f"Image not found: {image_path}")
+        if not os.path.exists(reference_image_path):
+            raise ValueError(f"Image not found: {reference_image_path}")
+        with open(image_path, "rb") as f:
+            image = f.read()
+        with open(reference_image_path, "rb") as f:
+            reference_image = f.read()
+        agent = GenerateImageInfer.build(
+            prompt=prompt,
+            action=Action.IMG2IMG,
+            sampler=Sampler.K_DPMPP_SDE,
+
+            image=image,
+            strength=0.6,
+
+            reference_image=reference_image,
+            reference_strength=0.6,
+            reference_information_extracted=1,
+
+            add_original_image=True,  # This Not affect the vibe generation
+            qualityToggle=True,
+        )
+        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, remove_sign=True
+        )
+    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.get_event_loop()
+loop.run_until_complete(generate())