From b689d8d569ba3a8c61515f8888908aa1b32686e2 Mon Sep 17 00:00:00 2001 From: sudoskys Date: Tue, 13 Feb 2024 14:08:21 +0800 Subject: [PATCH 1/3] :art: chore(server): add upscale endpoint Added the `/ai/upscale` endpoint to generate an image using the Upscale functionality. The generated image is returned as a zip file. Refactor the code in `server.py` to include the `/ai/upscale` endpoint and handle any exceptions that may occur. This commit also includes minor code formatting and linting changes. --- src/novelai_python/server.py | 30 +++++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/src/novelai_python/server.py b/src/novelai_python/server.py index 8eb15b6..fe82911 100644 --- a/src/novelai_python/server.py +++ b/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) From e69fd2c18d406a1ace872b122b9e78dba6a9a9db Mon Sep 17 00:00:00 2001 From: sudoskys Date: Tue, 13 Feb 2024 17:54:38 +0800 Subject: [PATCH 2/3] :sparkles: feat(generate_image): Add support for img2img action --- playground/generate_image_img2img.py | 4 ++-- src/novelai_python/sdk/ai/generate_image.py | 13 ++++++++----- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/playground/generate_image_img2img.py b/playground/generate_image_img2img.py index 150794c..68f89bf 100644 --- a/playground/generate_image_img2img.py +++ b/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 ) diff --git a/src/novelai_python/sdk/ai/generate_image.py b/src/novelai_python/sdk/ai/generate_image.py index 2fff470..11f933f 100644 --- a/src/novelai_python/sdk/ai/generate_image.py +++ b/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: From 5a7371ee602e575b322b4c42853bec7b2f856afe Mon Sep 17 00:00:00 2001 From: sudoskys Date: Tue, 13 Feb 2024 17:54:50 +0800 Subject: [PATCH 3/3] :sparkles: feat(generate_image): Add support for img2img action --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index c3af90f..25c1931 100644 --- a/pyproject.toml +++ b/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 = "coldlando@hotmail.com" },