Merge pull request #23 from LlmKira/dev

Dev

playground/generate_image_img2img.py

@@ -32,10 +32,10 @@ async def main():
     try:
         gen = GenerateImageInfer.build(
             prompt=f"1girl, spring, jacket, sfw, angel, flower,{enhance}",
-            action=Action.GENERATE,
+            action=Action.IMG2IMG,
             image=encoded,
             add_original_image=True,
-            strength=0.99,
+            strength=0.5,
             width=1088,
             height=896
         )

pyproject.toml

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

src/novelai_python/sdk/ai/generate_image.py

@@ -7,6 +7,7 @@
 import json
 import math
 import random
+from copy import deepcopy
 from enum import Enum, IntEnum
 from io import BytesIO
 from typing import Optional, Union
@@ -249,7 +250,7 @@ def model_post_init(self, *args) -> None:
     def validate_model(self):
         if self.action == Action.INFILL and not self.parameters.mask:
             logger.warning("Mask maybe required for infill mode.")
-        if self.action == Action.INFILL:
+        if self.action != Action.GENERATE:
             self.parameters.extra_noise_seed = self.parameters.seed
         return self
 
@@ -394,6 +395,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"
         }
 
     async def request(self,
@@ -421,20 +423,21 @@ async def request(self,
             session.headers.clear()
             session.headers.update(override_headers)
         try:
-            _log_data = request_data.copy()
+            _log_data = deepcopy(request_data)
             if self.action == Action.GENERATE:
                 logger.debug(f"Request Data: {_log_data}")
             else:
                 _log_data.get("parameters", {}).update({
-                    "image": "base64 data" if self.parameters.image else "None",
+                    "image": "base64 data" if self.parameters.image else None,
                 }
                 )
                 _log_data.get("parameters", {}).update(
                     {
-                        "mask": "base64 data" if self.parameters.mask else "None",
+                        "mask": "base64 data" if self.parameters.mask else None,
                     }
                 )
-                logger.debug(f"Request Data: {request_data}")
+                logger.debug(f"Request Data: {_log_data}")
+            del _log_data
         except Exception as e:
             logger.warning(f"Error when print log data: {e}")
         try:

src/novelai_python/server.py

@@ -15,6 +15,7 @@
 
 from .credential import JwtCredential, SecretStr
 from .sdk.ai.generate_image import GenerateImageInfer
+from .sdk.ai.upscale import Upscale
 from .sdk.user.information import Information
 from .sdk.user.login import Login
 from .sdk.user.subscription import Subscription
@@ -90,6 +91,33 @@ async def subscription(
         return JSONResponse(status_code=500, content=e.__dict__)
 
 
+@app.post("/ai/upscale")
+async def upscale(
+        req: Upscale,
+        current_token: str = Depends(get_current_token)
+):
+    """
+    生成图片
+    :param current_token: Authorization
+    :param req: Upscale
+    :return:
+    """
+    try:
+        _result = await req.request(session=get_session(current_token), remove_sign=True)
+        zip_file_bytes = io.BytesIO()
+        with zipfile.ZipFile(zip_file_bytes, mode="w", compression=zipfile.ZIP_DEFLATED) as zip_file:
+            file = _result.files  # ONLY TUPLE
+            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=image.zip'
+        })
+    except Exception as e:
+        logger.exception(e)
+        return JSONResponse(status_code=500, content=e.__dict__)
+
+
 @app.post("/ai/generate_image")
 async def generate_image(
         req: GenerateImageInfer,
@@ -110,7 +138,7 @@ async def generate_image(
         # 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'
+            'Content-Disposition': 'attachment;filename=image.zip'
         })
     except Exception as e:
         logger.exception(e)