From 51dcd09a09e9e4281455b7dd8125a780046fca0b Mon Sep 17 00:00:00 2001 From: Zenafey Date: Wed, 13 Mar 2024 21:34:39 +0500 Subject: [PATCH] 5.1 --- README.md | 5 +- prodiapy/__init__.py | 21 +- prodiapy/_client.py | 170 ++----- prodiapy/_exceptions.py | 63 ++- prodiapy/aio.py | 53 ++ prodiapy/resources/__init__.py | 13 +- prodiapy/resources/constants.py | 133 ----- prodiapy/resources/engine.py | 125 ++--- prodiapy/resources/facerestore.py | 61 +++ prodiapy/resources/faceswap.py | 108 ++-- prodiapy/resources/general.py | 286 +++++++---- prodiapy/resources/logger.py | 6 +- prodiapy/resources/response.py | 34 +- prodiapy/resources/stablediffusion.py | 505 +++++-------------- prodiapy/resources/stablediffusiongeneral.py | 406 +++++++++++++++ prodiapy/resources/stablediffusionxl.py | 303 +---------- prodiapy/resources/upscale.py | 118 +++-- prodiapy/resources/utils.py | 40 +- pyproject.toml | 8 +- 19 files changed, 1187 insertions(+), 1271 deletions(-) create mode 100644 prodiapy/aio.py delete mode 100644 prodiapy/resources/constants.py create mode 100644 prodiapy/resources/facerestore.py create mode 100644 prodiapy/resources/stablediffusiongeneral.py diff --git a/README.md b/README.md index 8acdb6a..9e40173 100644 --- a/README.md +++ b/README.md @@ -1,9 +1,12 @@ # Description prodiapy is an unofficial lightweight Python package for [Prodia Stable Diffusion API](https://docs.prodia.com/reference/getting-started) +# Pre-requirements +The minimal ver. of Python supported by this package is 3.8.10 + # Installation ```commandline -pip install -U prodiapy +pip install prodiapy ``` # text2img example diff --git a/prodiapy/__init__.py b/prodiapy/__init__.py index 6362564..df7248a 100644 --- a/prodiapy/__init__.py +++ b/prodiapy/__init__.py @@ -1,11 +1,10 @@ -from . import _exceptions -from . import resources -from .resources.logger import logger -from ._client import Prodia, AsyncProdia - -__all__ = [ - "_exceptions", - "resources", - "Prodia", - "AsyncProdia" -] +from . import _exceptions, resources, aio +from .resources.logger import logger +from ._client import Prodia + +__all__ = [ + "_exceptions", + "resources", + "Prodia", + "aio" +] diff --git a/prodiapy/_client.py b/prodiapy/_client.py index ef0ce49..e9c651e 100644 --- a/prodiapy/_client.py +++ b/prodiapy/_client.py @@ -1,115 +1,55 @@ -from prodiapy.resources.engine import SyncAPIClient, AsyncAPIClient -from prodiapy import resources -from prodiapy._exceptions import * -from typing_extensions import override - -import os - - -class Prodia(SyncAPIClient): - sd: resources.StableDiffusion - sdxl: resources.StableDiffusionXL - - api_key: str - - def __init__( - self, - api_key: str | None = None, - base_url: str | None = None - ) -> None: - """Construct a new prodia client instance. - - This automatically infers the following arguments from their corresponding environment variables if they are not provided: - - `api_key` from `PRODIA_API_KEY` - """ - if api_key is None: - api_key = os.environ.get("PRODIA_API_KEY") - if api_key is None: - raise AuthenticationError( - "The api_key client option must be set either by passing api_key to the client or by setting the \ - PRODIA_API_KEY environment variable" - ) - self.api_key = api_key - - if base_url is None: - base_url = f"https://api.prodia.com/v1" - - super().__init__( - base_url=base_url, - headers=self.auth_headers - ) - - self.sd = resources.StableDiffusion(self) - self.sdxl = resources.StableDiffusionXL(self) - - general = resources.General(self) - upscale = resources.Upscale(self) - faceswap = resources.FaceSwap(self) - - self.faceswap = faceswap.faceswap - self.upscale = upscale.upscale - self.create = general.create - self.job = general.job - self.constant = general.constant - self.wait = general.wait - - - @property - @override - def auth_headers(self) -> dict[str, str]: - api_key = self.api_key - return {'X-Prodia-Key': api_key, 'Content-Type': 'application/json'} - - -class AsyncProdia(AsyncAPIClient): - sd: resources.AsyncStableDiffusion - sdxl: resources.AsyncStableDiffusionXL - - api_key: str - - def __init__( - self, - api_key: str | None = None, - base_url: str | None = None - ) -> None: - """Construct a new async prodia client instance. - - This automatically infers the following arguments from their corresponding environment variables if they are not provided: - - `api_key` from `PRODIA_API_KEY` - """ - if api_key is None: - api_key = os.environ.get("PRODIA_API_KEY") - if api_key is None: - raise AuthenticationError( - "The api_key client option must be set either by passing api_key to the client or by setting the \ - PRODIA_API_KEY environment variable" - ) - self.api_key = api_key - - if base_url is None: - base_url = f"https://api.prodia.com/v1" - - super().__init__( - base_url=base_url, - headers=self.auth_headers - ) - - self.sd = resources.AsyncStableDiffusion(self) - self.sdxl = resources.AsyncStableDiffusionXL(self) - - general = resources.AsyncGeneral(self) - upscale = resources.AsyncUpscale(self) - faceswap = resources.AsyncFaceSwap(self) - - self.faceswap = faceswap.faceswap - self.upscale = upscale.upscale - self.create = general.create - self.job = general.job - self.constant = general.constant - self.wait = general.wait - - @property - @override - def auth_headers(self) -> dict[str, str]: - api_key = self.api_key - return {'X-Prodia-Key': api_key, 'Content-Type': 'application/json'} \ No newline at end of file + +from prodiapy.resources.engine import SyncAPIClient +from prodiapy import resources +from typing import Optional +from prodiapy._exceptions import * + +import os + + +class Prodia(SyncAPIClient): + api_key: str + sd: resources.StableDiffusion + sdxl: resources.StableDiffusionXL + general: resources.General + + def __init__( + self, + api_key: Optional[str] = os.getenv("PRODIA_API_KEY"), + base_url: str = os.getenv("PRODIA_API_BASE", "https://api.prodia.com/v1") + ): + """Construct a new prodia client instance. + + This automatically infers the following arguments from their corresponding environment variables if they are not provided: + - `api_key` from `PRODIA_API_KEY` + """ + if api_key is None: + raise AuthenticationError( + "The api_key client option must be set either by passing api_key to the client or by setting the \ +PRODIA_API_KEY environment variable" + ) + self.api_key = api_key + + super().__init__( + base_url=base_url, + headers=self.auth_headers + ) + + self.sd = resources.StableDiffusion(self) + self.sdxl = resources.StableDiffusionXL(self) + + general = resources.General(self) + + self.faceswap = general.faceswap + self.upscale = general.upscale + self.create = general.create + self.job = general.job + self.constant = general.constants + self.wait = general.wait + + @property + def auth_headers(self) -> dict: + api_key = self.api_key + return {'X-Prodia-Key': api_key, 'Content-Type': 'application/json'} + + diff --git a/prodiapy/_exceptions.py b/prodiapy/_exceptions.py index a251638..24e16da 100644 --- a/prodiapy/_exceptions.py +++ b/prodiapy/_exceptions.py @@ -1,28 +1,35 @@ - - -__all__ = [ - "AuthenticationError", - "InvalidParameterError", - "UnknownError", - "FailedJobError" -] - - -class ProdiaError(Exception): - pass - - -class AuthenticationError(ProdiaError): - pass - - -class InvalidParameterError(ProdiaError): - pass - - -class UnknownError(ProdiaError): - pass - - -class FailedJobError(ProdiaError): - pass + + +__all__ = [ + "AuthenticationError", + "InvalidParameterError", + "UnknownError", + "FailedJobError", + "exception_vocab" +] + + +class ProdiaError(Exception): + pass + + +class AuthenticationError(ProdiaError): + pass + + +class InvalidParameterError(ProdiaError): + pass + + +class UnknownError(ProdiaError): + pass + + +class FailedJobError(ProdiaError): + pass + + +exception_vocab = { + 401 | 401: AuthenticationError, + 400: InvalidParameterError + } diff --git a/prodiapy/aio.py b/prodiapy/aio.py new file mode 100644 index 0000000..50dd5cd --- /dev/null +++ b/prodiapy/aio.py @@ -0,0 +1,53 @@ + +from prodiapy.resources.engine import AsyncAPIClient +from prodiapy import resources +from typing import Optional +from prodiapy._exceptions import * + +import os + + +class Prodia(AsyncAPIClient): + api_key: str + sd: resources.AsyncStableDiffusion + sdxl: resources.AsyncStableDiffusionXL + general: resources.AsyncGeneral + + def __init__( + self, + api_key: Optional[str] = os.getenv("PRODIA_API_KEY"), + base_url: str = os.getenv("PRODIA_API_BASE", "https://api.prodia.com/v1") + ) -> None: + """Construct a new async prodia client instance. + + This automatically infers the following arguments from their corresponding environment variables if they are not provided: + - `api_key` from `PRODIA_API_KEY` + """ + if api_key is None: + raise AuthenticationError( + "The api_key client option must be set either by passing api_key to the client or by setting the \ +PRODIA_API_KEY environment variable" + ) + self.api_key = api_key + + super().__init__( + base_url=base_url, + headers=self.auth_headers + ) + + self.sd = resources.AsyncStableDiffusion(self) + self.sdxl = resources.AsyncStableDiffusionXL(self) + + general = resources.AsyncGeneral(self) + + self.faceswap = general.faceswap + self.upscale = general.upscale + self.create = general.create + self.job = general.job + self.constant = general.constants + self.wait = general.wait + + @property + def auth_headers(self) -> dict: + api_key = self.api_key + return {'X-Prodia-Key': api_key, 'Content-Type': 'application/json'} \ No newline at end of file diff --git a/prodiapy/resources/__init__.py b/prodiapy/resources/__init__.py index 9e25ad6..04e3ca8 100644 --- a/prodiapy/resources/__init__.py +++ b/prodiapy/resources/__init__.py @@ -1,6 +1,7 @@ -from .logger import logger -from .stablediffusion import StableDiffusion, AsyncStableDiffusion -from .stablediffusionxl import StableDiffusionXL, AsyncStableDiffusionXL -from .upscale import Upscale, AsyncUpscale -from .faceswap import FaceSwap, AsyncFaceSwap -from .general import General, AsyncGeneral +from .logger import logger +from .stablediffusion import StableDiffusion, AsyncStableDiffusion +from .stablediffusionxl import StableDiffusionXL, AsyncStableDiffusionXL +from .upscale import Upscale, AsyncUpscale +from .faceswap import FaceSwap, AsyncFaceSwap +from .facerestore import FaceRestore, AsyncFaceRestore +from .general import General, AsyncGeneral diff --git a/prodiapy/resources/constants.py b/prodiapy/resources/constants.py deleted file mode 100644 index 7315e80..0000000 --- a/prodiapy/resources/constants.py +++ /dev/null @@ -1,133 +0,0 @@ -from typing import Literal - - -sd_model_literal = Literal[ - "3Guofeng3_v34.safetensors [50f420de]", - "absolutereality_V16.safetensors [37db0fc3]", - "absolutereality_v181.safetensors [3d9d4d2b]", - "amIReal_V41.safetensors [0a8a2e61]", - "analog-diffusion-1.0.ckpt [9ca13f02]", - "anythingv3_0-pruned.ckpt [2700c435]", - "anything-v4.5-pruned.ckpt [65745d25]", - "anythingV5_PrtRE.safetensors [893e49b9]", - "AOM3A3_orangemixs.safetensors [9600da17]", - "blazing_drive_v10g.safetensors [ca1c1eab]", - "cetusMix_Version35.safetensors [de2f2560]", - "childrensStories_v13D.safetensors [9dfaabcb]", - "childrensStories_v1SemiReal.safetensors [a1c56dbb]", - "childrensStories_v1ToonAnime.safetensors [2ec7b88b]", - "Counterfeit_v30.safetensors [9e2a8f19]", - "cuteyukimixAdorable_midchapter3.safetensors [04bdffe6]", - "cyberrealistic_v33.safetensors [82b0d085]", - "dalcefo_v4.safetensors [425952fe]", - "deliberate_v2.safetensors [10ec4b29]", - "deliberate_v3.safetensors [afd9d2d4]", - "dreamlike-anime-1.0.safetensors [4520e090]", - "dreamlike-diffusion-1.0.safetensors [5c9fd6e0]", - "dreamlike-photoreal-2.0.safetensors [fdcf65e7]", - "dreamshaper_6BakedVae.safetensors [114c8abb]", - "dreamshaper_7.safetensors [5cf5ae06]", - "dreamshaper_8.safetensors [9d40847d]", - "edgeOfRealism_eorV20.safetensors [3ed5de15]", - "EimisAnimeDiffusion_V1.ckpt [4f828a15]", - "elldreths-vivid-mix.safetensors [342d9d26]", - "epicrealism_naturalSinRC1VAE.safetensors [90a4c676]", - "ICantBelieveItsNotPhotography_seco.safetensors [4e7a3dfd]", - "juggernaut_aftermath.safetensors [5e20c455]", - "lofi_v4.safetensors [ccc204d6]", - "lyriel_v16.safetensors [68fceea2]", - "majicmixRealistic_v4.safetensors [29d0de58]", - "mechamix_v10.safetensors [ee685731]", - "meinamix_meinaV9.safetensors [2ec66ab0]", - "meinamix_meinaV11.safetensors [b56ce717]", - "neverendingDream_v122.safetensors [f964ceeb]", - "openjourney_V4.ckpt [ca2f377f]", - "pastelMixStylizedAnime_pruned_fp16.safetensors [793a26e8]", - "portraitplus_V1.0.safetensors [1400e684]", - "protogenx34.safetensors [5896f8d5]", - "Realistic_Vision_V1.4-pruned-fp16.safetensors [8d21810b]", - "Realistic_Vision_V2.0.safetensors [79587710]", - "Realistic_Vision_V4.0.safetensors [29a7afaa]", - "Realistic_Vision_V5.0.safetensors [614d1063]", - "redshift_diffusion-V10.safetensors [1400e684]", - "revAnimated_v122.safetensors [3f4fefd9]", - "rundiffusionFX25D_v10.safetensors [cd12b0ee]", - "rundiffusionFX_v10.safetensors [cd4e694d]", - "sdv1_4.ckpt [7460a6fa]", - "v1-5-pruned-emaonly.safetensors [d7049739]", - "v1-5-inpainting.safetensors [21c7ab71]", - "shoninsBeautiful_v10.safetensors [25d8c546]", - "theallys-mix-ii-churned.safetensors [5d9225a4]", - "timeless-1.0.ckpt [7c4971d4]", - "toonyou_beta6.safetensors [980f6b15]" -] - -sdxl_model_literal = Literal[ - "dreamshaperXL10_alpha2.safetensors [c8afe2ef]", - "dynavisionXL_0411.safetensors [c39cc051]", - "juggernautXL_v45.safetensors [e75f5471]", - "realismEngineSDXL_v10.safetensors [af771c3f]", - "sd_xl_base_1.0.safetensors [be9edd61]" -] - -sd_sampler_literal = Literal[ - "Euler", - "Euler a", - "LMS", - "Heun", - "DPM2", - "DPM2 a", - "DPM++ 2S a", - "DPM++ 2M", - "DPM++ SDE", - "DPM fast", - "DPM adaptive", - "LMS Karras", - "DPM2 Karras", - "DPM2 a Karras", - "DPM++ 2S a Karras", - "DPM++ 2M Karras", - "DPM++ SDE Karras", - "DDIM", - "PLMS" -] - -sdxl_sampler_literal = Literal[ - "Euler", - "Euler a", - "LMS", - "Heun", - "DPM2", - "DPM2 a", - "DPM++ 2S a", - "DPM++ 2M", - "DPM++ SDE", - "DPM fast", - "DPM adaptive", - "LMS Karras", - "DPM2 Karras", - "DPM2 a Karras", - "DPM++ 2S a Karras", - "DPM++ 2M Karras", - "DPM++ SDE Karras" -] - -style_preset_literal = Literal[ - "3d-model", - "analog-film", - "anime", - "cinematic", - "comic-book", - "digital-art", - "enhance", - "fantasy-art", - "isometric", - "line-art", - "low-poly", - "neon-punk", - "origami", - "photographic", - "pixel-art", - "texture", - "craft-clay" -] \ No newline at end of file diff --git a/prodiapy/resources/engine.py b/prodiapy/resources/engine.py index 477016f..28fa0ca 100644 --- a/prodiapy/resources/engine.py +++ b/prodiapy/resources/engine.py @@ -1,73 +1,52 @@ -import aiohttp -import requests -from prodiapy._exceptions import * -from prodiapy.resources import logger - - -class SyncAPIClient: - def __init__(self, base_url, headers): - self.base_url = base_url - self.headers = headers - - def _request(self, method, endpoint, body=None): - r = requests.request(method, self.base_url+endpoint, json=body, headers=self.headers) - match r.status_code: - case 200: - return r.json() - case 401 | 402: - logger.error("Caught error(Unauthorized)") - raise AuthenticationError(f"Prodia API returned {r.status_code}. Details: {r.text}") - case 400: - logger.error("Caught error(Invalid Generation Parameters)") - raise InvalidParameterError(f"Prodia API returned {r.status_code}. Details: {r.text}") - case _: - logger.error("Unknown request error") - raise UnknownError(f"Prodia API returned {r.status_code}. Details: {r.text}") - - def _post(self, endpoint, body): - return self._request("post", endpoint, body) - - def _get(self, endpoint): - return self._request("get", endpoint) - - -class APIResource: - - def __init__(self, client) -> None: - self._client = client - self._get = client._get - self._post = client._post - - -class AsyncAPIClient: - base_url: str - headers: dict - - def __init__(self, base_url: str, headers: dict): - self.base_url = base_url - self.headers = headers - - async def _request(self, method, endpoint, body=None): - async with aiohttp.ClientSession() as s: - async with s.request(method, self.base_url+endpoint, json=body, headers=self.headers) as r: - match r.status: - case 200: - return await r.json() - case 401 | 402: - logger.error("Caught error(Unauthorized)") - raise AuthenticationError(f"Prodia API returned {r.status}. Details: {await r.text()}") - case 400: - logger.error("Caught error(Invalid Generation Parameters)") - raise InvalidParameterError(f"Prodia API returned {r.status}. Details: {await r.text()}") - case _: - logger.error("Unknown request error") - raise UnknownError(f"Prodia API returned {r.status}. Details: {await r.text()}") - - async def _post(self, endpoint, body): - return await self._request("post", endpoint, body) - - async def _get(self, endpoint): - return await self._request("get", endpoint) - - - +import aiohttp +import requests +from typing import Literal, Optional, Union +from prodiapy.resources.utils import raise_exception + + +class SyncAPIClient: + def __init__(self, base_url: str, headers: dict): + self.base_url = base_url + self.headers = headers + + def _request(self, method: Literal["get", "post"], endpoint: str, body: Optional[dict] = None): + r = requests.request(method, self.base_url+endpoint, json=body, headers=self.headers) + raise_exception(r.status_code, r.text) + + return r.json() + + def post(self, endpoint, body): return self._request("post", endpoint, body) + + def get(self, endpoint): return self._request("get", endpoint) + + +class AsyncAPIClient: + base_url: str + headers: dict + + def __init__(self, base_url: str, headers: dict): + self.base_url = base_url + self.headers = headers + + async def _request(self, method: Literal["get", "post"], endpoint: str, body: Optional[dict] = None): + async with aiohttp.ClientSession() as s: + async with s.request(method, self.base_url+endpoint, json=body, headers=self.headers) as r: + raise_exception(r.status, await r.text()) + + return await r.json() + + async def post(self, endpoint, body): return await self._request("post", endpoint, body) + + async def get(self, endpoint): return await self._request("get", endpoint) + + +class APIResource: + + def __init__(self, client: Union[SyncAPIClient, AsyncAPIClient]) -> None: + self._client = client + self._get = client.get + self._post = client.post + + + + diff --git a/prodiapy/resources/facerestore.py b/prodiapy/resources/facerestore.py new file mode 100644 index 0000000..d0b6de0 --- /dev/null +++ b/prodiapy/resources/facerestore.py @@ -0,0 +1,61 @@ + +from prodiapy.resources.engine import APIResource, SyncAPIClient, AsyncAPIClient +from prodiapy.resources.utils import form_body +from typing import Optional + + +class FaceRestore(APIResource): + def __init__(self, client: SyncAPIClient) -> None: + super().__init__(client) + + def facerestore( + self, + image_url: Optional[str] = None, + image_data: Optional[str] = None, + dict_parameters: Optional[dict] = None, + **kwargs + ) -> dict: + """ + Restore and enhance the face inside an image, source: https://docs.prodia.com/reference/facerestore + + Returns: + Python dictionary, containing job id + """ + return self._post( + "/facerestore", + body=form_body( + dict_parameters=dict_parameters, + imageUrl=image_url, + imageData=image_data, + **kwargs + ) + ) + + +class AsyncFaceRestore(APIResource): + def __init__(self, client: AsyncAPIClient) -> None: + super().__init__(client) + + async def facerestore( + self, + image_url: Optional[str] = None, + image_data: Optional[str] = None, + dict_parameters: Optional[dict] = None, + **kwargs + ) -> dict: + """ + Restore and enhance the face inside an image, source: https://docs.prodia.com/reference/facerestore + + Returns: + Python dictionary, containing job id + """ + return await self._post( + "/facerestore", + body=form_body( + dict_parameters=dict_parameters, + imageUrl=image_url, + imageData=image_data, + **kwargs + ) + ) + diff --git a/prodiapy/resources/faceswap.py b/prodiapy/resources/faceswap.py index 87886fb..f7e2b0a 100644 --- a/prodiapy/resources/faceswap.py +++ b/prodiapy/resources/faceswap.py @@ -1,43 +1,65 @@ -from prodiapy.resources.engine import APIResource -from prodiapy.resources.utils import form_body - - -class FaceSwap(APIResource): - def __init__(self, client) -> None: - super().__init__(client) - - def faceswap( - self, - sourceUrl: str | None = None, - targetUrl: str | None = None, - dict_parameters: dict | None = None - ) -> dict: - return self._post( - "/faceswap", - body=form_body( - dict_parameters=dict_parameters, - sourceUrl=sourceUrl, - targetUrl=targetUrl - ) - ) - - -class AsyncFaceSwap(APIResource): - def __init__(self, client) -> None: - super().__init__(client) - - async def faceswap( - self, - sourceUrl: str | None = None, - targetUrl: str | None = None, - dict_parameters: dict | None = None - ) -> dict: - return await self._post( - "/upscale", - body=form_body( - dict_parameters=dict_parameters, - sourceUrl=sourceUrl, - targetUrl=targetUrl - ) - ) - + +from prodiapy.resources.engine import APIResource, SyncAPIClient, AsyncAPIClient +from prodiapy.resources.utils import form_body +from typing import Optional + + +class FaceSwap(APIResource): + def __init__(self, client: SyncAPIClient) -> None: + super().__init__(client) + + def faceswap( + self, + source_url: Optional[str] = None, + target_url: Optional[str] = None, + dict_parameters: Optional[dict] = None, + **kwargs + ) -> dict: + """ + Swap a face inside an image (source_url) with another (target_url) +source: https://docs.prodia.com/reference/faceswap + + + Returns: + Python dictionary, containing job id + """ + return self._post( + "/faceswap", + body=form_body( + dict_parameters=dict_parameters, + sourceUrl=source_url, + targetUrl=target_url, + **kwargs + ) + ) + + +class AsyncFaceSwap(APIResource): + def __init__(self, client: AsyncAPIClient) -> None: + super().__init__(client) + + async def faceswap( + self, + source_url: Optional[str] = None, + target_url: Optional[str] = None, + dict_parameters: Optional[dict] = None, + **kwargs + ) -> dict: + """ + Swap a face inside an image (source_url) with another (target_url) +source: https://docs.prodia.com/reference/faceswap + + + Returns: + Python dictionary, containing job id + """ + return await self._post( + "/faceswap", + body=form_body( + dict_parameters=dict_parameters, + sourceUrl=source_url, + targetUrl=target_url, + **kwargs + ) + ) + diff --git a/prodiapy/resources/general.py b/prodiapy/resources/general.py index de094fa..5488048 100644 --- a/prodiapy/resources/general.py +++ b/prodiapy/resources/general.py @@ -1,102 +1,184 @@ -from prodiapy.resources.engine import APIResource -from prodiapy.resources.response import ProdiaResponse -from prodiapy.resources.logger import logger -from prodiapy.resources.utils import form_body -from prodiapy._exceptions import * - -import time -import asyncio - - -class General(APIResource): - def __init__(self, client) -> None: - super().__init__(client) - - def create( - self, - endpoint: str = "/sd/generate", - dict_parameters: dict | None = None, - **params - ) -> dict: - return self._post( - endpoint, - body=form_body( - dict_parameters=dict_parameters, - **params - ) - ) - - def constant( - self, - endpoint: str = "/sd/models" - ) -> list: - return self._get(endpoint) - - def job( - self, - job_id: str - ) -> dict: - return self._get(f"/job/{job_id}") - - def wait(self, job: dict, raise_on_fail: bool = True): - job_result = job - - while job_result['status'] not in ['succeeded', 'failed']: - time.sleep(0.25) - job_result = self.job(job['job']) - - if job_result['status'] == 'failed': - logger.failed(f"Job {job_result['job']} failed") - if raise_on_fail: - raise FailedJobError("Job failed") - return ProdiaResponse({'failed': True}) - - logger.success(f"Got result: {job_result}") - return ProdiaResponse(job_result) - - -class AsyncGeneral(APIResource): - def __init__(self, client) -> None: - super().__init__(client) - - async def create( - self, - endpoint: str = "/sd/generate", - dict_parameters: dict | None = None, - **params - ) -> dict: - return await self._post( - endpoint, - body=form_body( - dict_parameters=dict_parameters, - **params - ) - ) - - async def constant( - self, - endpoint: str = "/sd/models" - ) -> list: - return await self._get(endpoint) - - async def job( - self, - job_id: str - ) -> dict: - return await self._get(f"/job/{job_id}") - - async def wait(self, job: dict, raise_on_fail: bool = True): - job_result = job - - while job_result['status'] not in ['succeeded', 'failed']: - await asyncio.sleep(0.25) - job_result = await self.job(job['job']) - - if job_result['status'] == 'failed': - logger.failed(f"Job {job_result['job']} failed") - if raise_on_fail: - raise FailedJobError("Job failed") - return ProdiaResponse({'failed': True}) - - logger.success(f"Got result: {job_result}") - return ProdiaResponse(job_result) \ No newline at end of file + + +from prodiapy.resources.engine import APIResource, SyncAPIClient, AsyncAPIClient +from prodiapy.resources.facerestore import FaceRestore, AsyncFaceRestore +from prodiapy.resources.faceswap import FaceSwap, AsyncFaceSwap +from prodiapy.resources.upscale import Upscale, AsyncUpscale +from prodiapy.resources.response import ProdiaResponse +from prodiapy.resources.utils import form_body +from prodiapy.resources import logger +from prodiapy._exceptions import * + +import time +import asyncio +from typing import Optional + + +class General(APIResource): + facerestore: FaceRestore.facerestore + faceswap: FaceSwap.faceswap + upscale: Upscale.upscale + + def __init__(self, client: SyncAPIClient) -> None: + super().__init__(client) + self.facerestore = FaceRestore(client).facerestore + self.faceswap = FaceSwap(client).faceswap + self.upscale = Upscale(client).upscale + + def create( + self, + endpoint: str = "/sd/generate", + dict_parameters: Optional[dict] = None, + **params + ) -> dict: + """ + Universal method for any POST request to Prodia API + Args: + endpoint: endpoint in format "/{endpoint}" + dict_parameters: request body in dictionary format, if passed **params are ignored + **params: request body in kwargs format + + Returns: + Python dictionary, containing job id + """ + return self._post( + endpoint, + body=form_body( + dict_parameters=dict_parameters, + **params + ) + ) + + def constants(self, endpoint: str = "/sd/models") -> list: + """ + Universal endpoint for any GET request to Prodia API(excepting /job/{job_id}, use job() method for this) + Args: + endpoint: listing endpoint in format "/{endpoint}" + + Returns: + Actual list of available choices for chosen parameter + """ + return self._get(endpoint) + + def job(self, job_id: str) -> dict: + """ + Get job information, source: https://docs.prodia.com/reference/getjob + + Returns: + Python dictionary, containing job information(status, result if succeeded generation) + """ + return self._get(f"/job/{job_id}") + + def wait( + self, + job: dict, + raise_on_fail: bool = True + ) -> ProdiaResponse: + """ + Wait until generation complete and get result or error if generation failed + Args: + job: Python dictionary you got from POST methods + raise_on_fail: boolean + + Returns: + ProdiaResponse object + """ + job_result = job + + while job_result['status'] not in ['succeeded', 'failed']: + time.sleep(0.25) + job_result = self.job(job['job']) + + if job_result['status'] == 'failed': + logger.failed(f"Job {job_result['job']} failed") + if raise_on_fail: + raise FailedJobError(f"Job {job_result['job']} failed") + return ProdiaResponse({'job': job_result['job'], 'failed': True}) + + logger.success(f"Got result: {job_result}") + return ProdiaResponse(job_result) + + +class AsyncGeneral(APIResource): + facerestore: AsyncFaceRestore.facerestore + faceswap: AsyncFaceSwap.faceswap + upscale: AsyncUpscale.upscale + + def __init__(self, client: AsyncAPIClient) -> None: + super().__init__(client) + self.facerestore = AsyncFaceRestore(client).facerestore + self.faceswap = AsyncFaceSwap(client).faceswap + self.upscale = AsyncUpscale(client).upscale + + async def create( + self, + endpoint: str = "/sd/generate", + dict_parameters: Optional[dict] = None, + **params + ) -> dict: + """ + Universal method for any POST request to Prodia API + Args: + endpoint: endpoint in format "/{endpoint}" + dict_parameters: request body in dictionary format, if passed **params are ignored + **params: request body in kwargs format + + Returns: + Python dictionary, containing job id + """ + return await self._post( + endpoint, + body=form_body( + dict_parameters=dict_parameters, + **params + ) + ) + + async def constants(self, endpoint: str = "/sd/models") -> list: + """ + Universal endpoint for any GET request to Prodia API(excepting /job/{job_id}, use job() method for this) + Args: + endpoint: listing endpoint in format "/{endpoint}" + + Returns: + Actual list of available choices for chosen parameter + """ + return await self._get(endpoint) + + async def job(self, job_id: str) -> dict: + """ + Get information about a generation job, including status, source: https://docs.prodia.com/reference/getjob + + Returns: + Python dictionary, containing job information(status, result if succeeded generation) + """ + return await self._get(f"/job/{job_id}") + + async def wait( + self, + job: dict, + raise_on_fail: bool = True + ) -> ProdiaResponse: + """ + Wait until generation complete and get result or error if generation failed + Args: + job: Python dictionary you got from POST methods + raise_on_fail: boolean + + Returns: + ProdiaResponse object + """ + job_result = job + + while job_result['status'] not in ['succeeded', 'failed']: + await asyncio.sleep(0.25) + job_result = await self.job(job['job']) + + if job_result['status'] == 'failed': + logger.failed(f"Job {job_result['job']} failed") + if raise_on_fail: + raise FailedJobError(f"Job {job_result['job']} failed") + return ProdiaResponse({'job': job_result['job'], 'failed': True}) + + logger.success(f"Got result: {job_result}") + return ProdiaResponse(job_result) diff --git a/prodiapy/resources/logger.py b/prodiapy/resources/logger.py index 329215b..24cf196 100644 --- a/prodiapy/resources/logger.py +++ b/prodiapy/resources/logger.py @@ -1,3 +1,3 @@ -from zenlogger import Logger - -logger = Logger("prodiapy") +from zenlogger import Logger + +logger = Logger("prodiapy") diff --git a/prodiapy/resources/response.py b/prodiapy/resources/response.py index 347e3d2..f626768 100644 --- a/prodiapy/resources/response.py +++ b/prodiapy/resources/response.py @@ -1,17 +1,17 @@ - - -class ProdiaResponse: - """ - Prodia API response(after waiting for result) - - Attributes: - job_id: id of the job - image_url: URL of generated image - failed: is job failed or not(boolean) - json: JSON response(raw) - """ - def __init__(self, output: dict): - self.job_id = output.get('job') - self.image_url = output.get('imageUrl') - self.failed = output.get('failed', False) - self.json = output + + +class ProdiaResponse: + """ + Prodia API result class + + Attributes: + job_id: id of the job + image_url: URL of generated image + failed: is job failed or not(boolean) + json: JSON response(raw) + """ + def __init__(self, output: dict): + self.job_id = output.get('job') + self.image_url = output.get('imageUrl') + self.failed = output.get('failed', False) + self.json = output diff --git a/prodiapy/resources/stablediffusion.py b/prodiapy/resources/stablediffusion.py index 25a32a0..e16969a 100644 --- a/prodiapy/resources/stablediffusion.py +++ b/prodiapy/resources/stablediffusion.py @@ -1,387 +1,118 @@ -from prodiapy.resources.engine import APIResource -from typing import Union -from prodiapy.resources.constants import * -from prodiapy.resources.utils import form_body - - -class StableDiffusion(APIResource): - def __init__(self, client) -> None: - super().__init__(client) - - def generate( - self, - model: Union[str, sd_model_literal, None] = None, - prompt: str | None = None, - negative_prompt: str | None = None, - style_preset: Union[str, style_preset_literal, None] = None, - steps: int | None = None, - cfg_scale: int | float | None = None, - seed: int | None = None, - upscale: bool | None = None, - sampler: Union[str, sd_sampler_literal, None] = None, - aspect_ratio: Union[str, Literal["square", "portrait", "landscape"], None] = None, - width: int | None = None, - height: int | None = None, - dict_parameters: dict | None = None - ) -> dict: - return self._post( - "/sd/generate", - body=form_body( - dict_parameters=dict_parameters, - model=model, - prompt=prompt, - negative_prompt=negative_prompt, - style_preset=style_preset, - steps=steps, - cfg_scale=cfg_scale, - seed=seed, - upscale=upscale, - sampler=sampler, - aspect_ratio=aspect_ratio, - width=width, - height=height - ) - ) - - def transform( - self, - imageUrl: str | None = None, - imageData = None, - model: Union[str, sd_model_literal, None] = None, - prompt: str | None = None, - denoising_strength: float | None = None, - negative_prompt: str | None = None, - style_preset: Union[str, style_preset_literal, None] = None, - steps: int | None = None, - cfg_scale: int | float | None = None, - seed: int | None = None, - upscale: bool | None = None, - sampler: Union[str, sd_sampler_literal, None] = None, - width: int | None = None, - height: int | None = None, - dict_parameters: dict | None = None - ) -> dict: - return self._post( - "/sd/transform", - body=form_body( - dict_parameters=dict_parameters, - imageUrl=imageUrl, - imageData=imageData, - model=model, - prompt=prompt, - denoising_strength=denoising_strength, - negative_prompt=negative_prompt, - style_preset=style_preset, - steps=steps, - cfg_scale=cfg_scale, - seed=seed, - upscale=upscale, - sampler=sampler, - width=width, - height=height - ) - ) - - def inpainting( - self, - imageUrl: str | None = None, - imageData = None, - maskUrl: str | None = None, - maskData=None, - model: Union[str, sd_model_literal, None] = None, - prompt: str | None = None, - denoising_strength: float | None = None, - negative_prompt: str | None = None, - style_preset: Union[str, style_preset_literal, None] = None, - steps: int | None = None, - cfg_scale: int | float | None = None, - seed: int | None = None, - upscale: bool | None = None, - mask_blur: int | None = None, - inpainting_fill: Union[int, Literal[0, 1, 2, 3], None] = None, - inpainting_mask_invert: Union[int, Literal[0, 1], None] = None, - inpainting_full_res: bool | None = None, - sampler: Union[str, sd_sampler_literal, None] = None, - width: int | None = None, - height: int | None = None, - dict_parameters: dict | None = None - ) -> dict: - return self._post( - "/sd/inpainting", - body=form_body( - dict_parameters=dict_parameters, - imageUrl=imageUrl, - imageData=imageData, - maskUrl=maskUrl, - maskData=maskData, - model=model, - prompt=prompt, - denoising_strength=denoising_strength, - negative_prompt=negative_prompt, - style_preset=style_preset, - steps=steps, - cfg_scale=cfg_scale, - seed=seed, - upscale=upscale, - mask_blur=mask_blur, - inpainting_fill=inpainting_fill, - inpainting_mask_invert=inpainting_mask_invert, - inpainting_full_res=inpainting_full_res, - sampler=sampler, - width=width, - height=height - ) - ) - - def controlnet( - self, - imageUrl: str | None = None, - imageData = None, - model: Union[str, sd_model_literal, None] = None, - controlnet_model: str | None = None, - controlnet_module: str | None = None, - controlnet_mode: Union[int, Literal[0, 1 , 2], None] = None, - threshold_a: int | None = None, - threshold_b: int | None = None, - resize_mode: Union[int, Literal[0, 1, 2], None] = None, - prompt: str | None = None, - negative_prompt: str | None = None, - style_preset: Union[str, style_preset_literal, None] = None, - steps: int | None = None, - cfg_scale: int | float | None = None, - seed: int | None = None, - upscale: bool | None = None, - sampler: Union[str, sd_sampler_literal, None] = None, - width: int | None = None, - height: int | None = None, - dict_parameters: dict | None = None - ) -> dict: - return self._post( - "/sd/controlnet", - body=form_body( - dict_parameters=dict_parameters, - imageUrl=imageUrl, - imageData=imageData, - model=model, - controlnet_model=controlnet_model, - controlnet_module=controlnet_module, - controlnet_mode=controlnet_mode, - threshold_a=threshold_a, - threshold_b=threshold_b, - resize_mode=resize_mode, - prompt=prompt, - negative_prompt=negative_prompt, - style_preset=style_preset, - steps=steps, - cfg_scale=cfg_scale, - seed=seed, - upscale=upscale, - sampler=sampler, - width=width, - height=height - ) - ) - - def models(self) -> list: - return self._get("/sd/models") - - def samplers(self) -> list: - return self._get("/sd/samplers") - - def loras(self) -> list: - return self._get("/sd/loras") - - def embeddings(self) -> list: - return self._get("/sd/embeddings") - - -class AsyncStableDiffusion(APIResource): - def __init__(self, client) -> None: - super().__init__(client) - - async def generate( - self, - model: Union[str, sd_model_literal, None] = None, - prompt: str | None = None, - negative_prompt: str | None = None, - style_preset: Union[str, style_preset_literal, None] = None, - steps: int | None = None, - cfg_scale: int | float | None = None, - seed: int | None = None, - upscale: bool | None = None, - sampler: Union[str, sd_sampler_literal, None] = None, - aspect_ratio: Union[str, Literal["square", "portrait", "landscape"], None] = None, - width: int | None = None, - height: int | None = None, - dict_parameters: dict | None = None - ) -> dict: - return await self._post( - "/sd/generate", - body=form_body( - dict_parameters=dict_parameters, - model=model, - prompt=prompt, - negative_prompt=negative_prompt, - style_preset=style_preset, - steps=steps, - cfg_scale=cfg_scale, - seed=seed, - upscale=upscale, - sampler=sampler, - aspect_ratio=aspect_ratio, - width=width, - height=height - ) - ) - - async def transform( - self, - imageUrl: str | None = None, - imageData=None, - model: Union[str, sd_model_literal, None] = None, - prompt: str | None = None, - denoising_strength: float | None = None, - negative_prompt: str | None = None, - style_preset: Union[str, style_preset_literal, None] = None, - steps: int | None = None, - cfg_scale: int | float | None = None, - seed: int | None = None, - upscale: bool | None = None, - sampler: Union[str, sd_sampler_literal, None] = None, - width: int | None = None, - height: int | None = None, - dict_parameters: dict | None = None - ) -> dict: - return await self._post( - "/sd/transform", - body=form_body( - dict_parameters=dict_parameters, - imageUrl=imageUrl, - imageData=imageData, - model=model, - prompt=prompt, - denoising_strength=denoising_strength, - negative_prompt=negative_prompt, - style_preset=style_preset, - steps=steps, - cfg_scale=cfg_scale, - seed=seed, - upscale=upscale, - sampler=sampler, - width=width, - height=height - ) - ) - - async def inpainting( - self, - imageUrl: str | None = None, - imageData=None, - maskUrl: str | None = None, - maskData=None, - model: Union[str, sd_model_literal, None] = None, - prompt: str | None = None, - denoising_strength: float | None = None, - negative_prompt: str | None = None, - style_preset: Union[str, style_preset_literal, None] = None, - steps: int | None = None, - cfg_scale: int | float | None = None, - seed: int | None = None, - upscale: bool | None = None, - mask_blur: int | None = None, - inpainting_fill: Union[int, Literal[0, 1, 2, 3], None] = None, - inpainting_mask_invert: Union[int, Literal[0, 1], None] = None, - inpainting_full_res: bool | None = None, - sampler: Union[str, sd_sampler_literal, None] = None, - width: int | None = None, - height: int | None = None, - dict_parameters: dict | None = None - ) -> dict: - return await self._post( - "/sd/inpainting", - body=form_body( - dict_parameters=dict_parameters, - imageUrl=imageUrl, - imageData=imageData, - maskUrl=maskUrl, - maskData=maskData, - model=model, - prompt=prompt, - denoising_strength=denoising_strength, - negative_prompt=negative_prompt, - style_preset=style_preset, - steps=steps, - cfg_scale=cfg_scale, - seed=seed, - upscale=upscale, - mask_blur=mask_blur, - inpainting_fill=inpainting_fill, - inpainting_mask_invert=inpainting_mask_invert, - inpainting_full_res=inpainting_full_res, - sampler=sampler, - width=width, - height=height - ) - ) - - async def controlnet( - self, - imageUrl: str | None = None, - imageData=None, - model: Union[str, sd_model_literal, None] = None, - controlnet_model: str | None = None, - controlnet_module: str | None = None, - controlnet_mode: Union[int, Literal[0, 1, 2], None] = None, - threshold_a: int | None = None, - threshold_b: int | None = None, - resize_mode: Union[int, Literal[0, 1, 2], None] = None, - prompt: str | None = None, - negative_prompt: str | None = None, - style_preset: Union[str, style_preset_literal, None] = None, - steps: int | None = None, - cfg_scale: int | float | None = None, - seed: int | None = None, - upscale: bool | None = None, - sampler: Union[str, sd_sampler_literal, None] = None, - width: int | None = None, - height: int | None = None, - dict_parameters: dict | None = None - ) -> dict: - return await self._post( - "/sd/controlnet", - body=form_body( - dict_parameters=dict_parameters, - imageUrl=imageUrl, - imageData=imageData, - model=model, - controlnet_model=controlnet_model, - controlnet_module=controlnet_module, - controlnet_mode=controlnet_mode, - threshold_a=threshold_a, - threshold_b=threshold_b, - resize_mode=resize_mode, - prompt=prompt, - negative_prompt=negative_prompt, - style_preset=style_preset, - steps=steps, - cfg_scale=cfg_scale, - seed=seed, - upscale=upscale, - sampler=sampler, - width=width, - height=height - ) - ) - - async def models(self) -> list: - return await self._get("/sd/models") - - async def samplers(self) -> list: - return await self._get("/sd/samplers") - - async def loras(self) -> list: - return await self._get("/sd/loras") - - async def embeddings(self) -> list: - return await self._get("/sd/embeddings") - +from __future__ import annotations + +from prodiapy.resources.engine import APIResource +from typing import Union, Optional, Literal +from prodiapy.resources.stablediffusiongeneral import StableDiffusionGeneral, AsyncStableDiffusionGeneral +from prodiapy.resources.engine import SyncAPIClient, AsyncAPIClient +from prodiapy.resources.utils import form_body + + +class StableDiffusion(StableDiffusionGeneral): + """class related to /sd endpoints, source: https://docs.prodia.com/reference/generate""" + def __init__(self, client: SyncAPIClient) -> None: + super().__init__(client, model_architecture="sd") + + def controlnet( + self, + image_url: Optional[str] = None, + image_data: Optional[str] = None, + model: Optional[str] = None, + controlnet_model: Optional[str] = None, + controlnet_module: Optional[str] = None, + controlnet_mode: Optional[Union[int, Literal[0, 1, 2]]] = None, + threshold_a: Optional[int] = None, + threshold_b: Optional[int] = None, + resize_mode: Optional[Union[int, Literal[0, 1, 2]]] = None, + prompt: Optional[str] = None, + negative_prompt: Optional[str] = None, + style_preset: Optional[str] = None, + steps: Optional[int] = None, + cfg_scale: Optional[int | float] = None, + seed: Optional[int] = None, + sampler: Optional[str] = None, + width: Optional[int] = None, + height: Optional[int] = None, + dict_parameters: Optional[dict] = None, + **kwargs + ) -> dict: + return self._post( + f"/sd/controlnet", + body=form_body( + dict_parameters=dict_parameters, + imageUrl=image_url, + imageData=image_data, + model=model, + controlnet_model=controlnet_model, + controlnet_module=controlnet_module, + controlnet_mode=controlnet_mode, + threshold_a=threshold_a, + threshold_b=threshold_b, + resize_mode=resize_mode, + prompt=prompt, + negative_prompt=negative_prompt, + style_preset=style_preset, + steps=steps, + cfg_scale=cfg_scale, + seed=seed, + sampler=sampler, + width=width, + height=height, + **kwargs + ) + ) + + +class AsyncStableDiffusion(AsyncStableDiffusionGeneral): + """class related to /sd endpoints, source: https://docs.prodia.com/reference/generate""" + def __init__(self, client: AsyncAPIClient) -> None: + super().__init__(client, model_architecture="sd") + + async def controlnet( + self, + image_url: Optional[str] = None, + image_data: Optional[str] = None, + model: Optional[str] = None, + controlnet_model: Optional[str] = None, + controlnet_module: Optional[str] = None, + controlnet_mode: Optional[Union[int, Literal[0, 1, 2]]] = None, + threshold_a: Optional[int] = None, + threshold_b: Optional[int] = None, + resize_mode: Optional[Union[int, Literal[0, 1, 2]]] = None, + prompt: Optional[str] = None, + negative_prompt: Optional[str] = None, + style_preset: Optional[str] = None, + steps: Optional[int] = None, + cfg_scale: Optional[int | float] = None, + seed: Optional[int] = None, + sampler: Optional[str] = None, + width: Optional[int] = None, + height: Optional[int] = None, + dict_parameters: Optional[dict] = None, + **kwargs + ) -> dict: + return await self._post( + f"/sd/controlnet", + body=form_body( + dict_parameters=dict_parameters, + imageUrl=image_url, + imageData=image_data, + model=model, + controlnet_model=controlnet_model, + controlnet_module=controlnet_module, + controlnet_mode=controlnet_mode, + threshold_a=threshold_a, + threshold_b=threshold_b, + resize_mode=resize_mode, + prompt=prompt, + negative_prompt=negative_prompt, + style_preset=style_preset, + steps=steps, + cfg_scale=cfg_scale, + seed=seed, + sampler=sampler, + width=width, + height=height, + **kwargs + ) + ) + diff --git a/prodiapy/resources/stablediffusiongeneral.py b/prodiapy/resources/stablediffusiongeneral.py new file mode 100644 index 0000000..6d1e149 --- /dev/null +++ b/prodiapy/resources/stablediffusiongeneral.py @@ -0,0 +1,406 @@ +from __future__ import annotations + +from prodiapy.resources.engine import APIResource +from typing import Union, Optional, Literal +from prodiapy.resources.utils import form_body + + +class StableDiffusionGeneral(APIResource): + def __init__(self, client, model_architecture="sd") -> None: + super().__init__(client) + self.model_architecture = model_architecture + + def generate( + self, + model: Optional[str] = None, + prompt: Optional[str] = None, + negative_prompt: Optional[str] = None, + style_preset: Optional[str] = None, + steps: Optional[int] = None, + cfg_scale: Optional[int | float] = None, + seed: Optional[int] = None, + sampler: Optional[str] = None, + width: Optional[int] = None, + height: Optional[int] = None, + dict_parameters: Optional[dict] = None, + **kwargs + ) -> dict: + """ + text2image generation, sources: +https://docs.prodia.com/reference/generate for StableDiffusion +https://docs.prodia.com/reference/sdxl-generate for StableDiffusionXL + + Returns: + Python dictionary containing job id + + """ + return self._post( + f"/{self.model_architecture}/generate", + body=form_body( + dict_parameters=dict_parameters, + model=model, + prompt=prompt, + negative_prompt=negative_prompt, + style_preset=style_preset, + steps=steps, + cfg_scale=cfg_scale, + seed=seed, + sampler=sampler, + width=width, + height=height, + **kwargs + ) + ) + + def transform( + self, + image_url: Optional[str] = None, + image_data: Optional[str] = None, + model: Optional[str] = None, + prompt: Optional[str] = None, + denoising_strength: Optional[float] = None, + negative_prompt: Optional[str] = None, + style_preset: Optional[str] = None, + steps: Optional[int] = None, + cfg_scale: Optional[int | float] = None, + seed: Optional[int] = None, + sampler: Optional[str] = None, + width: Optional[int] = None, + height: Optional[int] = None, + dict_parameters: Optional[dict] = None, + **kwargs + ) -> dict: + """ + image2image generation, sources: +https://docs.prodia.com/reference/transform for StableDiffusion +https://docs.prodia.com/reference/sdxl-transform for StableDiffusionXL + + Returns: + Python dictionary containing job id + + """ + return self._post( + f"/{self.model_architecture}/transform", + body=form_body( + dict_parameters=dict_parameters, + imageUrl=image_url, + imageData=image_data, + model=model, + prompt=prompt, + denoising_strength=denoising_strength, + negative_prompt=negative_prompt, + style_preset=style_preset, + steps=steps, + cfg_scale=cfg_scale, + seed=seed, + sampler=sampler, + width=width, + height=height, + **kwargs + ) + ) + + def inpainting( + self, + image_url: Optional[str] = None, + image_data: Optional[str] = None, + mask_url: Optional[str] = None, + mask_data: Optional[str] = None, + model: Optional[str] = None, + prompt: Optional[str] = None, + denoising_strength: Optional[float] = None, + negative_prompt: Optional[str] = None, + style_preset: Optional[str] = None, + steps: Optional[int] = None, + cfg_scale: Optional[int | float] = None, + seed: Optional[int] = None, + mask_blur: Optional[int] = None, + inpainting_fill: Optional[Union[int, Literal[0, 1, 2, 3]]] = None, + inpainting_mask_invert: Optional[Union[int, Literal[0, 1]]] = None, + inpainting_full_res: Optional[bool] = None, + sampler: Optional[str] = None, + width: Optional[int] = None, + height: Optional[int] = None, + dict_parameters: Optional[dict] = None, + **kwargs + ) -> dict: + """ + Controlled image2image generation, sources: +https://docs.prodia.com/reference/inpainting for StableDiffusion +https://docs.prodia.com/reference/sdxl-inpainting for StableDiffusionXL + + Returns: + Python dictionary containing job id + + """ + return self._post( + f"/{self.model_architecture}/inpainting", + body=form_body( + dict_parameters=dict_parameters, + imageUrl=image_url, + imageData=image_data, + maskUrl=mask_url, + maskData=mask_data, + model=model, + prompt=prompt, + denoising_strength=denoising_strength, + negative_prompt=negative_prompt, + style_preset=style_preset, + steps=steps, + cfg_scale=cfg_scale, + seed=seed, + mask_blur=mask_blur, + inpainting_fill=inpainting_fill, + inpainting_mask_invert=inpainting_mask_invert, + inpainting_full_res=inpainting_full_res, + sampler=sampler, + width=width, + height=height, + **kwargs + ) + ) + + def models(self) -> list: + """ + Actual list of available models, sources: +https://docs.prodia.com/reference/listmodels for StableDiffusion +https://docs.prodia.com/reference/listsdxlmodels for StableDiffusionXL + + Returns: + Python list containing string names of available models + """ + return self._get(f"/{self.model_architecture}/models") + + def samplers(self) -> list: + """ + Actual list of available samplers, sources: +https://docs.prodia.com/reference/listsamplers for StableDiffusion +https://docs.prodia.com/reference/listsdxlsamplers for StableDiffusionXL + + Returns: + Python list containing string names of available samplers + """ + return self._get(f"/{self.model_architecture}/samplers") + + def loras(self) -> list: + """ + Actual list of available LoRa models, sources: +https://docs.prodia.com/reference/listloras for StableDiffusion +https://docs.prodia.com/reference/listsdxlloras for StableDiffusionXL + + Returns: + Python list containing string names of available LoRa models + """ + return self._get(f"/{self.model_architecture}/loras") + + def embeddings(self) -> list: + """ + Actual list of available embedding models, sources: +https://docs.prodia.com/reference/listembeddings for StableDiffusion +https://docs.prodia.com/reference/listsdxlembeddings for StableDiffusionXL + + Returns: + Python list containing string names of available embedding models + """ + return self._get(f"/{self.model_architecture}/embeddings") + + +class AsyncStableDiffusionGeneral(APIResource): + def __init__(self, client, model_architecture="sd") -> None: + super().__init__(client) + self.model_architecture = model_architecture + + async def generate( + self, + model: Optional[str] = None, + prompt: Optional[str] = None, + negative_prompt: Optional[str] = None, + style_preset: Optional[str] = None, + steps: Optional[int] = None, + cfg_scale: Optional[int | float] = None, + seed: Optional[int] = None, + sampler: Optional[str] = None, + width: Optional[int] = None, + height: Optional[int] = None, + dict_parameters: Optional[dict] = None, + **kwargs + ) -> dict: + """ + text2image generation, sources: +https://docs.prodia.com/reference/generate for StableDiffusion +https://docs.prodia.com/reference/sdxl-generate for StableDiffusionXL + + Returns: + Python dictionary containing job id + + """ + return await self._post( + f"/{self.model_architecture}/generate", + body=form_body( + dict_parameters=dict_parameters, + model=model, + prompt=prompt, + negative_prompt=negative_prompt, + style_preset=style_preset, + steps=steps, + cfg_scale=cfg_scale, + seed=seed, + sampler=sampler, + width=width, + height=height, + **kwargs + ) + ) + + async def transform( + self, + image_url: Optional[str] = None, + image_data: Optional[str] = None, + model: Optional[str] = None, + prompt: Optional[str] = None, + denoising_strength: Optional[float] = None, + negative_prompt: Optional[str] = None, + style_preset: Optional[str] = None, + steps: Optional[int] = None, + cfg_scale: Optional[int | float] = None, + seed: Optional[int] = None, + sampler: Optional[str] = None, + width: Optional[int] = None, + height: Optional[int] = None, + dict_parameters: Optional[dict] = None, + **kwargs + ) -> dict: + """ + image2image generation, sources: +https://docs.prodia.com/reference/transform for StableDiffusion +https://docs.prodia.com/reference/sdxl-transform for StableDiffusionXL + + Returns: + Python dictionary containing job id + + """ + return await self._post( + f"/{self.model_architecture}/transform", + body=form_body( + dict_parameters=dict_parameters, + imageUrl=image_url, + imageData=image_data, + model=model, + prompt=prompt, + denoising_strength=denoising_strength, + negative_prompt=negative_prompt, + style_preset=style_preset, + steps=steps, + cfg_scale=cfg_scale, + seed=seed, + sampler=sampler, + width=width, + height=height, + **kwargs + ) + ) + + async def inpainting( + self, + image_url: Optional[str] = None, + image_data: Optional[str] = None, + mask_url: Optional[str] = None, + mask_data: Optional[str] = None, + model: Optional[str] = None, + prompt: Optional[str] = None, + denoising_strength: Optional[float] = None, + negative_prompt: Optional[str] = None, + style_preset: Optional[str] = None, + steps: Optional[int] = None, + cfg_scale: Optional[int | float] = None, + seed: Optional[int] = None, + mask_blur: Optional[int] = None, + inpainting_fill: Optional[Union[int, Literal[0, 1, 2, 3]]] = None, + inpainting_mask_invert: Optional[Union[int, Literal[0, 1]]] = None, + inpainting_full_res: Optional[bool] = None, + sampler: Optional[str] = None, + width: Optional[int] = None, + height: Optional[int] = None, + dict_parameters: Optional[dict] = None, + **kwargs + ) -> dict: + """ + Controlled image2image generation, sources: +https://docs.prodia.com/reference/inpainting for StableDiffusion +https://docs.prodia.com/reference/sdxl-inpainting for StableDiffusionXL + + Returns: + Python dictionary containing job id + + """ + return await self._post( + f"/{self.model_architecture}/inpainting", + body=form_body( + dict_parameters=dict_parameters, + imageUrl=image_url, + imageData=image_data, + maskUrl=mask_url, + maskData=mask_data, + model=model, + prompt=prompt, + denoising_strength=denoising_strength, + negative_prompt=negative_prompt, + style_preset=style_preset, + steps=steps, + cfg_scale=cfg_scale, + seed=seed, + mask_blur=mask_blur, + inpainting_fill=inpainting_fill, + inpainting_mask_invert=inpainting_mask_invert, + inpainting_full_res=inpainting_full_res, + sampler=sampler, + width=width, + height=height, + **kwargs + ) + ) + + async def models(self) -> list: + """ + Actual list of available models, sources: +https://docs.prodia.com/reference/listmodels for StableDiffusion +https://docs.prodia.com/reference/listsdxlmodels for StableDiffusionXL + + Returns: + Python list containing string names of available models + """ + return await self._get(f"/{self.model_architecture}/models") + + async def samplers(self) -> list: + """ + Actual list of available samplers, sources: +https://docs.prodia.com/reference/listsamplers for StableDiffusion +https://docs.prodia.com/reference/listsdxlsamplers for StableDiffusionXL + + Returns: + Python list containing string names of available samplers + """ + return await self._get(f"/{self.model_architecture}/samplers") + + async def loras(self) -> list: + """ + Actual list of available LoRa models, sources: +https://docs.prodia.com/reference/listloras for StableDiffusion +https://docs.prodia.com/reference/listsdxlloras for StableDiffusionXL + + Returns: + Python list containing string names of available LoRa models + """ + return await self._get(f"/{self.model_architecture}/loras") + + async def embeddings(self) -> list: + """ + Actual list of available embedding models, sources: +https://docs.prodia.com/reference/listembeddings for StableDiffusion +https://docs.prodia.com/reference/listsdxlembeddings for StableDiffusionXL + + Returns: + Python list containing string names of available embedding models + """ + return await self._get(f"/{self.model_architecture}/embeddings") + diff --git a/prodiapy/resources/stablediffusionxl.py b/prodiapy/resources/stablediffusionxl.py index 57cd813..be998be 100644 --- a/prodiapy/resources/stablediffusionxl.py +++ b/prodiapy/resources/stablediffusionxl.py @@ -1,285 +1,18 @@ -from prodiapy.resources.engine import APIResource -from typing import Union -from prodiapy.resources.constants import * -from prodiapy.resources.utils import form_body - - -class StableDiffusionXL(APIResource): - def __init__(self, client) -> None: - super().__init__(client) - - def generate( - self, - model: Union[str, sdxl_model_literal, None] = None, - prompt: str | None = None, - negative_prompt: str | None = None, - style_preset: Union[str, style_preset_literal, None] = None, - steps: int | None = None, - cfg_scale: int | float | None = None, - seed: int | None = None, - upscale: bool | None = None, - sampler: Union[str, sdxl_sampler_literal, None] = None, - width: int | None = None, - height: int | None = None, - dict_parameters: dict | None = None - ) -> dict: - return self._post( - "/sdxl/generate", - body=form_body( - dict_parameters=dict_parameters, - model=model, - prompt=prompt, - negative_prompt=negative_prompt, - style_preset=style_preset, - steps=steps, - cfg_scale=cfg_scale, - seed=seed, - upscale=upscale, - sampler=sampler, - width=width, - height=height - ) - ) - - def transform( - self, - imageUrl: str | None = None, - imageData = None, - model: Union[str, sdxl_model_literal, None] = None, - prompt: str | None = None, - denoising_strength: float | None = None, - negative_prompt: str | None = None, - style_preset: Union[str, style_preset_literal, None] = None, - steps: int | None = None, - cfg_scale: int | float | None = None, - seed: int | None = None, - upscale: bool | None = None, - sampler: Union[str, sdxl_sampler_literal, None] = None, - width: int | None = None, - height: int | None = None, - dict_parameters: dict | None = None - ) -> dict: - return self._post( - "/sdxl/transform", - body=form_body( - dict_parameters=dict_parameters, - imageUrl=imageUrl, - imageData=imageData, - model=model, - prompt=prompt, - denoising_strength=denoising_strength, - negative_prompt=negative_prompt, - style_preset=style_preset, - steps=steps, - cfg_scale=cfg_scale, - seed=seed, - upscale=upscale, - sampler=sampler, - width=width, - height=height - ) - ) - - def inpainting( - self, - imageUrl: str | None = None, - imageData = None, - maskUrl: str | None = None, - maskData=None, - model: Union[str, sdxl_model_literal, None] = None, - prompt: str | None = None, - denoising_strength: float | None = None, - negative_prompt: str | None = None, - style_preset: Union[str, style_preset_literal, None] = None, - steps: int | None = None, - cfg_scale: int | float | None = None, - seed: int | None = None, - upscale: bool | None = None, - mask_blur: int | None = None, - inpainting_fill: Union[int, Literal[0, 1, 2, 3], None] = None, - inpainting_mask_invert: Union[int, Literal[0, 1], None] = None, - inpainting_full_res: bool | None = None, - sampler: Union[str, sdxl_sampler_literal, None] = None, - width: int | None = None, - height: int | None = None, - dict_parameters: dict | None = None - ) -> dict: - return self._post( - "/sdxl/inpainting", - body=form_body( - dict_parameters=dict_parameters, - imageUrl=imageUrl, - imageData=imageData, - maskUrl=maskUrl, - maskData=maskData, - model=model, - prompt=prompt, - denoising_strength=denoising_strength, - negative_prompt=negative_prompt, - style_preset=style_preset, - steps=steps, - cfg_scale=cfg_scale, - seed=seed, - upscale=upscale, - mask_blur=mask_blur, - inpainting_fill=inpainting_fill, - inpainting_mask_invert=inpainting_mask_invert, - inpainting_full_res=inpainting_full_res, - sampler=sampler, - width=width, - height=height - ) - ) - - def models(self) -> list: - return self._get("/sdxl/models") - - def samplers(self) -> list: - return self._get("/sdxl/samplers") - - def loras(self) -> list: - return self._get("/sdxl/loras") - - def embeddings(self) -> list: - return self._get("/sdxl/embeddings") - - -class AsyncStableDiffusionXL(APIResource): - def __init__(self, client) -> None: - super().__init__(client) - - async def generate( - self, - model: Union[str, sdxl_model_literal, None] = None, - prompt: str | None = None, - negative_prompt: str | None = None, - style_preset: Union[str, style_preset_literal, None] = None, - steps: int | None = None, - cfg_scale: int | float | None = None, - seed: int | None = None, - upscale: bool | None = None, - sampler: Union[str, sdxl_sampler_literal, None] = None, - width: int | None = None, - height: int | None = None, - dict_parameters: dict | None = None - ) -> dict: - return await self._post( - "/sdxl/generate", - body=form_body( - dict_parameters=dict_parameters, - model=model, - prompt=prompt, - negative_prompt=negative_prompt, - style_preset=style_preset, - steps=steps, - cfg_scale=cfg_scale, - seed=seed, - upscale=upscale, - sampler=sampler, - width=width, - height=height - ) - ) - - async def transform( - self, - imageUrl: str | None = None, - imageData = None, - model: Union[str, sdxl_model_literal, None] = None, - prompt: str | None = None, - denoising_strength: float | None = None, - negative_prompt: str | None = None, - style_preset: Union[str, style_preset_literal, None] = None, - steps: int | None = None, - cfg_scale: int | float | None = None, - seed: int | None = None, - upscale: bool | None = None, - sampler: Union[str, sdxl_sampler_literal, None] = None, - width: int | None = None, - height: int | None = None, - dict_parameters: dict | None = None - ) -> dict: - return await self._post( - "/sdxl/transform", - body=form_body( - dict_parameters=dict_parameters, - imageUrl=imageUrl, - imageData=imageData, - model=model, - prompt=prompt, - denoising_strength=denoising_strength, - negative_prompt=negative_prompt, - style_preset=style_preset, - steps=steps, - cfg_scale=cfg_scale, - seed=seed, - upscale=upscale, - sampler=sampler, - width=width, - height=height - ) - ) - - async def inpainting( - self, - imageUrl: str | None = None, - imageData = None, - maskUrl: str | None = None, - maskData=None, - model: Union[str, sdxl_model_literal, None] = None, - prompt: str | None = None, - denoising_strength: float | None = None, - negative_prompt: str | None = None, - style_preset: Union[str, style_preset_literal, None] = None, - steps: int | None = None, - cfg_scale: int | float | None = None, - seed: int | None = None, - upscale: bool | None = None, - mask_blur: int | None = None, - inpainting_fill: Union[int, Literal[0, 1, 2, 3], None] = None, - inpainting_mask_invert: Union[int, Literal[0, 1], None] = None, - inpainting_full_res: bool | None = None, - sampler: Union[str, sdxl_sampler_literal, None] = None, - width: int | None = None, - height: int | None = None, - dict_parameters: dict | None = None - ) -> dict: - return await self._post( - "/sdxl/inpainting", - body=form_body( - dict_parameters=dict_parameters, - imageUrl=imageUrl, - imageData=imageData, - maskUrl=maskUrl, - maskData=maskData, - model=model, - prompt=prompt, - denoising_strength=denoising_strength, - negative_prompt=negative_prompt, - style_preset=style_preset, - steps=steps, - cfg_scale=cfg_scale, - seed=seed, - upscale=upscale, - mask_blur=mask_blur, - inpainting_fill=inpainting_fill, - inpainting_mask_invert=inpainting_mask_invert, - inpainting_full_res=inpainting_full_res, - sampler=sampler, - width=width, - height=height - ) - ) - - async def models(self) -> list: - return await self._get("/sdxl/models") - - async def samplers(self) -> list: - return await self._get("/sdxl/samplers") - - async def loras(self) -> list: - return await self._get("/sdxl/loras") - - async def embeddings(self) -> list: - return await self._get("/sdxl/embeddings") - + +from prodiapy.resources.stablediffusiongeneral import StableDiffusionGeneral, AsyncStableDiffusionGeneral + + +class StableDiffusionXL(StableDiffusionGeneral): + """ + class related to /sdxl endpoints, source: https://docs.prodia.com/reference/sdxl-generate + """ + def __init__(self, client) -> None: + super().__init__(client, model_architecture="sdxl") + + +class AsyncStableDiffusionXL(AsyncStableDiffusionGeneral): + """ + class related to /sdxl endpoints, source: https://docs.prodia.com/reference/sdxl-generate + """ + def __init__(self, client) -> None: + super().__init__(client, model_architecture="sdxl") diff --git a/prodiapy/resources/upscale.py b/prodiapy/resources/upscale.py index 80fbded..8e4f076 100644 --- a/prodiapy/resources/upscale.py +++ b/prodiapy/resources/upscale.py @@ -1,49 +1,69 @@ -from prodiapy.resources.engine import APIResource -from typing import Union -from prodiapy.resources.constants import * -from prodiapy.resources.utils import form_body - - -class Upscale(APIResource): - def __init__(self, client) -> None: - super().__init__(client) - - def upscale( - self, - imageUrl: str | None = None, - imageData=None, - resize: Union[int, Literal[2, 4]] = 2, - dict_parameters: dict | None = None - ) -> dict: - return self._post( - "/upscale", - body=form_body( - dict_parameters=dict_parameters, - imageUrl=imageUrl, - imageData=imageData, - resize=resize - ) - ) - - -class AsyncUpscale(APIResource): - def __init__(self, client) -> None: - super().__init__(client) - - async def upscale( - self, - imageUrl: str | None = None, - imageData=None, - resize: Union[int, Literal[2, 4], None] = None, - dict_parameters: dict | None = None - ) -> dict: - return await self._post( - "/upscale", - body=form_body( - dict_parameters=dict_parameters, - imageUrl=imageUrl, - imageData=imageData, - resize=resize - ) - ) - + +from prodiapy.resources.engine import APIResource, SyncAPIClient, AsyncAPIClient +from typing import Union, Literal, Optional +from prodiapy.resources.utils import form_body + + +class Upscale(APIResource): + def __init__(self, client: SyncAPIClient) -> None: + super().__init__(client) + + def upscale( + self, + image_url: Optional[str] = None, + image_data: Optional[str] = None, + resize: Optional[Union[int, Literal[2, 4]]] = None, + model: Optional[str] = None, + dict_parameters: Optional[dict] = None, + **kwargs + ) -> dict: + """ + image upscaling, source: https://docs.prodia.com/reference/upscale + + Returns: + Python dictionary containing job id + """ + return self._post( + "/upscale", + body=form_body( + dict_parameters=dict_parameters, + imageUrl=image_url, + imageData=image_data, + resize=resize, + model=model, + **kwargs + ) + ) + + +class AsyncUpscale(APIResource): + def __init__(self, client: AsyncAPIClient) -> None: + super().__init__(client) + + async def upscale( + self, + image_url: Optional[str] = None, + image_data: Optional[str] = None, + resize: Optional[Union[int, Literal[2, 4]]] = None, + model: Optional[str] = None, + dict_parameters: Optional[dict] = None, + **kwargs + ) -> dict: + """ + image upscaling, source: https://docs.prodia.com/reference/upscale + + Returns: + Python dictionary containing job id + """ + return await self._post( + "/upscale", + body=form_body( + dict_parameters=dict_parameters, + imageUrl=image_url, + imageData=image_data, + resize=resize, + model=model, + **kwargs + ) + ) + diff --git a/prodiapy/resources/utils.py b/prodiapy/resources/utils.py index 223c6c2..e40f65c 100644 --- a/prodiapy/resources/utils.py +++ b/prodiapy/resources/utils.py @@ -1,14 +1,26 @@ - - -def form_body(dict_parameters: dict | None = None, **kwargs): - match dict_parameters: - case None: - body = {} - for kwarg in kwargs: - if kwargs.get(kwarg) is not None: - body[kwarg] = kwargs.get(kwarg) - case _: - body = dict_parameters - - return body - +from typing import Optional +from prodiapy._exceptions import * + + +def form_body(dict_parameters: Optional[dict] = None, **kwargs): + if dict_parameters: + body = dict_parameters + else: + body = {} + for kwarg in kwargs: + if kwargs.get(kwarg) is not None: + body[kwarg] = kwargs.get(kwarg) + + return body + + +def raise_exception(status: int, message: str) -> None: + message_body = f"Prodia API returned {status}. Details: {message}" + if status == 200: + pass + elif status in exception_vocab: + raise exception_vocab[status](message_body) + else: + raise UnknownError(message_body) + + diff --git a/pyproject.toml b/pyproject.toml index 2cb88aa..cb516aa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,14 +1,14 @@ [tool.poetry] name = "prodiapy" -version = "4.6" +version = "5.1" description = "Package for using Prodia API" authors = ["zenafey "] readme = "README.md" [tool.poetry.dependencies] -python = "^3.11" -requests = "^2.27.1" -aiohttp = "^3.8.1" +python = "^3.8" +requests = "^2.31.0" +aiohttp = "^3.9.1" zenlogger = "^0.0.2" [tool.poetry.dev-dependencies]