Python SDK for Stable Diffusion API (Txt2Img/Img2Img/ControlNet/VAE)
# Configuration file for the Sphinx documentation builder.
#
# For the full list of built-in configuration values, see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

# -- Project information -----------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information

import os
import sys

sys.path.insert(0, os.path.abspath('..'))

project = 'novita-python-sdk'
copyright = '2023, novita'
author = 'novita'
release = '0.0.1'

# -- General configuration ---------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

extensions = [
    'sphinx.ext.autodoc',
    'sphinx.ext.viewcode',
    'sphinx.ext.napoleon'
]

templates_path = ['_templates']
exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']


# -- Options for HTML output -------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output

html_theme = 'alabaster'
html_static_path = ['_static']

docs/conf.py
novitalabs/python-sdk
#!/usr/bin/env python
# -*- coding: UTF-8 -*-

import os

from novita_client import *

# get your api key refer to https://docs.novita.ai/get-started/
client = NovitaClient(os.getenv('NOVITA_API_KEY'))

controlnet_model = client.models().filter_by_type(ModelType.CONTROLNET).get_by_name("control_v1p_sd15_qrcode_monster_v2")
if controlnet_model is None:
    raise Exception("controlnet model not found")

req = Txt2ImgRequest(
    prompt="a beautify butterfly in the colorful flowers, best quality, best details, masterpiece",
    sampler_name=Samplers.DPMPP_M_KARRAS,
    width=512,
    height=512,
    steps=30,
    controlnet_units=[
        ControlnetUnit(
            input_image=read_image_to_base64(os.path.join(os.path.abspath(os.path.dirname(__file__)), "fixtures/qrcode.png")),
            control_mode=ControlNetMode.BALANCED,
            model=controlnet_model.sd_name,
            module=ControlNetPreprocessor.NULL,
            resize_mode=ControlNetResizeMode.JUST_RESIZE,
            weight=2.0,
        )
    ]
)

res = client.sync_txt2img(req)
if res.data.status != ProgressResponseStatusCode.SUCCESSFUL:
    raise Exception('Failed to generate image with error: ' +
                    res.data.failed_reason)

save_image(res.data.imgs_bytes[0], "qrcode-art.png")

examples/controlnet_qrcode.py
novitalabs/python-sdk
#!/usr/bin/env python
# -*- coding: UTF-8 -*-

from novita_client import NovitaClient, ModelType
# get your api key refer to https://docs.novita.ai/get-started/
client = NovitaClient(os.getenv('NOVITA_API_KEY'))

# filter by model type
print("lora count", len(client.models().filter_by_type(ModelType.LORA)))
print("checkpoint count", len(client.models().filter_by_type(ModelType.CHECKPOINT)))
print("textinversion count", len(
    client.models().filter_by_type(ModelType.TEXT_INVERSION)))
print("vae count", len(client.models().filter_by_type(ModelType.VAE)))
print("controlnet count", len(client.models().filter_by_type(ModelType.CONTROLNET)))


# filter by civitai tags
client.models().filter_by_civi_tags('anime')

# filter by nsfw
client.models().filter_by_nsfw(False)  # or True

# sort by civitai download
client.models().sort_by_civitai_download()

# chain filters
client.models().\
    filter_by_type(ModelType.CHECKPOINT).\
    filter_by_nsfw(False).\
    filter_by_civitai_tags('anime')

examples/model_search.py
novitalabs/python-sdk
import os

from novita_client import *

client = NovitaClient(os.getenv('NOVITA_API_KEY'))
req = Txt2ImgRequest(
    model_name='dreamshaper_8_93211.safetensors',
    prompt='a dog flying in the sky',
    width=512,
    height=512,
    batch_size=1,
    cfg_scale=7.5,
    sampler_name=Samplers.EULER_A,
    enable_hr=True,
    hr_scale=2.0
)

res = client.sync_txt2img(req)
if res.data.status != ProgressResponseStatusCode.SUCCESSFUL:
    raise Exception('Failed to generate image with error: ' +
                    res.data.failed_reason)

save_image(res.data.imgs_bytes[0], "txt2img-hiresfix-1024.png")

examples/txt2img_with_hiresfix.py
novitalabs/python-sdk
#!/usr/bin/env python
# -*- coding: UTF-8 -*-

import os
from novita_client import NovitaClient, Txt2ImgRequest, Samplers, ProgressResponseStatusCode, ModelType, add_lora_to_prompt, save_image


client = NovitaClient(os.getenv('NOVITA_API_KEY'))

res = client.sync_txt2img(Txt2ImgRequest(
    prompt="modern american comic about superman vs batman, digital color comicbook style, Batman jumps out of the way as Superman charges towards him, but Superman quickly recovers and delivers a powerful kick",
    batch_size=1,
    cfg_scale=7.0,
    sampler_name=Samplers.DPMPP_M_KARRAS,
    model_name="sd_xl_base_1.0.safetensors",
    height=1024,
    width=1024,
    seed=2563190745,
))

if res.data.status != ProgressResponseStatusCode.SUCCESSFUL:
    raise Exception('Failed to generate image with error: ' +
                    res.data.failed_reason)
save_image(res.data.imgs_bytes[0], "test.png")

examples/txt2img_with_lora.py
novitalabs/python-sdk
import os

from novita_client import *

client = NovitaClient(os.getenv('NOVITA_API_KEY'))
req = Txt2ImgRequest(
    model_name='sd_xl_base_1.0.safetensors',
    prompt='a dog flying in the sky',
    width=1024,
    height=1024,
    batch_size=1,
    cfg_scale=7.5,
    sampler_name=Samplers.EULER_A,
    sd_refiner=Refiner(
        checkpoint='sd_xl_refiner_1.0.safetensors',
        switch_at=0.5,
    ))

res = client.sync_txt2img(req)
if res.data.status != ProgressResponseStatusCode.SUCCESSFUL:
    raise Exception('Failed to generate image with error: ' +
                    res.data.failed_reason)

save_image(res.data.imgs_bytes[0], "txt2img-refiner.png")

examples/txt2img_with_refiner.py
novitalabs/python-sdk
#!/usr/bin/env python
# -*- coding: UTF-8 -*-

from .novita import *
from .proto import *
from .utils import *

src/novita_client/__init__.py
novitalabs/python-sdk
#!/usr/bin/env python
# -*- coding: UTF-8 -*-

class NovitaError(Exception):
    pass


class NovitaResponseError(NovitaError):
    pass


class NovitaTimeoutError(NovitaError):
    pass

src/novita_client/exceptions.py
novitalabs/python-sdk
#!/usr/bin/env python
# -*- coding: UTF-8 -*-

import logging

from time import sleep

from .version import __version__

from .exceptions import *
from .proto import *

import requests
from . import settings


logger = logging.getLogger(__name__)


class NovitaClient:
    """NovitaClient is the main entry point for interacting with the Novita API."""

    def __init__(self, api_key):
        self.base_url = "https://api.novita.ai/v2"
        self.api_key = api_key
        self.session = requests.Session()

        if not self.api_key:
            raise ValueError("NOVITA_API_KEY environment variable not set")

        # eg: {"all": [proto.ModelInfo], "checkpoint": [proto.ModelInfo], "lora": [proto.ModelInfo]}
        self._model_list_cache = None
        self._extra_headers = {}

    def set_extra_headers(self, headers: dict):
        self._extra_headers = headers

    def _get(self, api_path, params=None) -> dict:
        headers = {
            'Accept': 'application/json',
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}',
            'User-Agent': "novita-python-sdk/{}".format(__version__),
            'Accept-Encoding': 'gzip, deflate',
        }
        headers.update(self._extra_headers)

        logger.debug(f"[GET] params: {params}")

        response = self.session.get(
            self.base_url + api_path,
            headers=headers,
            params=params,
            timeout=settings.DEFAULT_REQUEST_TIMEOUT,
        )

        logger.debug(f"[GET] response: {response.content}")
        if response.status_code != 200:
            logger.error(f"Request failed: {response}")
            raise NovitaResponseError(
                f"Request failed with status {response.status_code}")

        return response.json()

    def _post(self, api_path, data) -> dict:
        headers = {
            'Accept': 'application/json',
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}',
            'User-Agent': "novita-python-sdk/{}".format(__version__),
            'Accept-Encoding': 'gzip, deflate',
        }
        headers.update(self._extra_headers)

        logger.debug(f"[POST] data: {data}")

        response = self.session.post(
            self.base_url + api_path,
            headers=headers,
            json=data,
            timeout=settings.DEFAULT_REQUEST_TIMEOUT,
        )

        logger.debug(f"[POST] response: {response.content}")
        if response.status_code != 200:
            logger.error(f"Request failed: {response}")
            raise NovitaResponseError(
                f"Request failed with status {response.status_code}")

        return response.json()

    def txt2img(self, request: Txt2ImgRequest) -> Txt2ImgResponse:
        """Asynchronously generate images from request

        Args:
            request (Txt2ImgRequest): The request object containing the text and image generation parameters.

        Returns:
            Txt2ImgResponse: The response object containing the task ID and status URL.
        """
        response = self._post('/txt2img', request.to_dict())

        return Txt2ImgResponse.from_dict(response)

    def progress(self, task_id: str) -> ProgressResponse:
        """Get the progress of a task.

        Args:
            task_id (str): The ID of the task to get the progress for.

        Returns:
            ProgressResponse: The response object containing the progress information for the task.
        """
        response = self._get('/progress', {
            'task_id': task_id,
        })

        return ProgressResponse.from_dict(response)

    def img2img(self, request: Img2ImgRequest) -> Img2ImgResponse:
        """Asynchronously generate images from request

        Args:
            request (Img2ImgRequest): The request object containing the image and image generation parameters.

        Returns:
            Img2ImgResponse: The response object containing the task ID and status URL.
        """
        response = self._post('/img2img', request.to_dict())

        return Img2ImgResponse.from_dict(response)

    def wait_for_task(self, task_id, wait_for: int = 300, callback: callable = None) -> ProgressResponse:
        """Wait for a task to complete

        This method waits for a task to complete by periodically checking its progress. If the task is not completed within the specified time, an NovitaTimeoutError is raised.

        Args:
            task_id (_type_): The ID of the task to wait for.
            wait_for (int, optional): The maximum time to wait for the task to complete, in seconds. Defaults to 300.

        Raises:
            NovitaTimeoutError: If the task fails to complete within the specified time.

        Returns:
            ProgressResponse: The response object containing the progress information for the task.
        """
        i = 0

        while i < wait_for:
            logger.info(f"Waiting for task {task_id} to complete")

            progress = self.progress(task_id)

            if callback and callable(callback):
                try:
                    callback(progress)
                except Exception as e:
                    logger.error(f"Task {task_id} progress callback failed: {e}")

            logger.info(
                f"Task {task_id} progress eta_relative: {progress.data.eta_relative}")

            if progress.data.status.finished():
                logger.info(f"Task {task_id} completed")
                return progress

            sleep(settings.DEFAULT_POLL_INTERVAL)
            i += 1

        raise NovitaTimeoutError(
            f"Task {task_id} failed to complete in {wait_for} seconds")

    def sync_txt2img(self, request: Txt2ImgRequest, download_images=True, callback: callable = None) -> ProgressResponse:
        """Synchronously generate images from request, optionally download images

        This method generates images synchronously from the given request object. If download_images is set to True, the generated images will be downloaded.

        Args:
            request (Txt2ImgRequest): The request object containing the input text and other parameters.
            download_images (bool, optional): Whether to download the generated images. Defaults to True.

        Raises:
            NovitaResponseError: If the text to image generation fails.

        Returns:
            ProgressResponse: The response object containing the task status and generated images.
        """
        response = self.txt2img(request)

        if response.data is None:
            raise NovitaResponseError(f"Text to Image generation failed with response {response.msg}, code: {response.code}")

        res = self.wait_for_task(response.data.task_id, callback=callback)
        if download_images:
            res.download_images()
        return res

    def sync_img2img(self, request: Img2ImgRequest, download_images=True, callback: callable = None) -> ProgressResponse:
        """Synchronously generate images from request, optionally download images

        Args:
            request (Img2ImgRequest): The request object containing the input image and other parameters.
            download_images (bool, optional): Whether to download the generated images. Defaults to True.

        Returns:
            ProgressResponse: The response object containing the task status and generated images.
        """
        response = self.img2img(request)

        if response.data is None:
            raise NovitaResponseError(f"Image to Image generation failed with response {response.msg}, code: {response.code}")

        res = self.wait_for_task(response.data.task_id, callback=callback)
        if download_images:
            res.download_images()
        return res

    def sync_upscale(self, request: UpscaleRequest, download_images=True, callback: callable = None) -> ProgressResponse:
        """Syncronously upscale image from request, optionally download images

        Args:
            request (UpscaleRequest): _description_
            download_images (bool, optional): _description_. Defaults to True.

        Returns:
            ProgressResponse: _description_
        """
        response = self.upscale(request)

        if response.data is None:
            raise NovitaResponseError(f"Upscale failed with response {response.msg}, code: {response.code}")

        res = self.wait_for_task(response.data.task_id, callback=callback)
        if download_images:
            res.download_images()
        return res

    def upscale(self, request: UpscaleRequest) -> UpscaleResponse:
        """Upscale image

        This method sends a request to the Novita API to upscale an image using the specified parameters.

        Args:
            request (UpscaleRequest): An object containing the input image and other parameters.

        Returns:
            UpscaleResponse: An object containing the task status and the URL of the upscaled image.
        """
        response = self._post('/upscale', request.to_dict())

        return UpscaleResponse.from_dict(response)

    def models(self, refresh=False) -> ModelList:
        """Get list of models

        This method retrieves a list of models available in the Novita API. If the list has already been retrieved and
        `refresh` is False, the cached list will be returned. Otherwise, a new request will be made to the API to
        retrieve the list.

        Args:
            refresh (bool, optional): If True, a new request will be made to the API to retrieve the list of models.
                If False and the list has already been retrieved, the cached list will be returned. Defaults to False.

        Returns:
            ModelList: A list of models available in the Novita API.
        """

        if (self._model_list_cache is None or len(self._model_list_cache) == 0) or refresh:
            res = self._get('/models')

            # TODO: fix this
            res_controlnet = self._get(
                '/models', params={'type': 'controlnet'})
            res_vae = self._get('/models', params={'type': 'vae'})

            tmp = []
            tmp.extend(MoodelsResponse.from_dict(res).data.models)
            tmp.extend(MoodelsResponse.from_dict(res_controlnet).data.models)
            tmp.extend(MoodelsResponse.from_dict(res_vae).data.models)

            # In future /models maybe return all models, so we need to filter out duplicates
            tmp_set = set()
            models = []
            for m in tmp:
                if m.sd_name not in tmp_set:
                    tmp_set.add(m.sd_name)
                    models.append(m)

            self._model_list_cache = ModelList(models)

        return self._model_list_cache

src/novita_client/novita.py
novitalabs/python-sdk
#!/usr/bin/env python
# -*- coding: UTF-8 -*-

from dataclasses import dataclass, field
from typing import Optional, List
from .serializer import JSONe
from enum import Enum
from .utils import batch_download_images


# --------------- ControlNet ---------------

class ControlNetMode(Enum):
    BALANCED = 0
    PROMPT_IMPORTANCE = 1
    CONTROLNET_IMPORTANCE = 2

    def __str__(self):
        return self.name


class ControlNetResizeMode(Enum):
    JUST_RESIZE = 0
    RESIZE_OR_CORP = 1
    RESIZE_AND_FILL = 2

    def __str__(self):
        return self.name


class ControlNetPreprocessor(Enum):
    NULL = 'none'
    CANNY = 'canny'
    DEPTH = 'depth'
    DEPTH_LERES = 'depth_leres'
    DEPTH_LERES_PLUS_PLUS = 'depth_leres++'
    HED = 'hed'
    HED_SAFE = 'hed_safe'
    MEDIAPIPE_FACE = 'mediapipe_face'
    MLSD = 'mlsd'
    NORMAL_MAP = 'normal_map'
    OPENPOSE = 'openpose'
    OPENPOSE_HAND = 'openpose_hand'
    OPENPOSE_FACE = 'openpose_face'
    OPENPOSE_FACEONLY = 'openpose_faceonly'
    OPENPOSE_FULL = 'openpose_full'
    CLIP_VISION = 'clip_vision'
    COLOR = 'color'
    PIDINET = 'pidinet'
    PIDINET_SAFE = 'pidinet_safe'
    PIDINET_SKETCH = 'pidinet_sketch'
    PIDINET_SCRIBBLE = 'pidinet_scribble'
    SCRIBBLE_XDOG = 'scribble_xdog'
    SCRIBBLE_HED = 'scribble_hed'
    SEGMENTATION = 'segmentation'
    THRESHOLD = 'threshold'
    DEPTH_ZOE = 'depth_zoe'
    NORMAL_BAE = 'normal_bae'
    ONEFORMER_COCO = 'oneformer_coco'
    ONEFORMER_ADE20K = 'oneformer_ade20k'
    LINEART = 'lineart'
    LINEART_COARSE = 'lineart_coarse'
    LINEART_ANIME = 'lineart_anime'
    LINEART_STANDARD = 'lineart_standard'
    SHUFFLE = 'shuffle'
    TILE_RESAMPLE = 'tile_resample'
    INVERT = 'invert'
    LINEART_ANIME_DENOISE = 'lineart_anime_denoise'
    REFERENCE_ONLY = 'reference_only'
    REFERENCE_ADAIN = 'reference_adain'
    REFERENCE_ADAIN_PLUS_ATTN = 'reference_adain+attn'
    INPAINT = 'inpaint'
    INPAINT_ONLY = 'inpaint_only'
    INPAINT_ONLY_PLUS_LAMA = 'inpaint_only+lama'
    TILE_COLORFIX = 'tile_colorfix'
    TILE_COLORFIX_PLUS_SHARP = 'tile_colorfix+sharp'

    def __str__(self):
        return self.name


@dataclass
class ControlnetUnit(JSONe):
    model: str
    weight: Optional[float] = 1
    module: Optional[ControlNetPreprocessor] = ControlNetPreprocessor.NULL
    input_image: Optional[str] = None
    control_mode: Optional[ControlNetMode] = ControlNetMode.BALANCED
    resize_mode: Optional[ControlNetResizeMode] = ControlNetResizeMode.RESIZE_OR_CORP
    mask: Optional[str] = None
    processor_res: Optional[int] = 512
    threshold_a: Optional[int] = 64
    threshold_b: Optional[int] = 64
    guidance_start: Optional[float] = 0.0
    guidance_end: Optional[float] = 1.0
    pixel_perfect: Optional[bool] = False


# --------------- Samplers ---------------
@dataclass
class Samplers:
    EULER_A = 'Euler a'
    EULER = 'Euler'
    LMS = 'LMS'
    HEUN = 'Heun'
    DPM2 = 'DPM2'
    DPM2_A = 'DPM2 a'
    DPM2_KARRAS = 'DPM2 Karras'
    DPM2_A_KARRAS = 'DPM2 a Karras'
    DPMPP_S_A = 'DPM++ 2S a'
    DPMPP_M = 'DPM++ 2M'
    DPMPP_SDE = 'DPM++ SDE'
    DPMPP_KARRAS = 'DPM++ Karras'
    DPMPP_S_A_KARRAS = 'DPM++ 2S a Karras'
    DPMPP_M_KARRAS = 'DPM++ 2M Karras'
    DPMPP_SDE_KARRAS = 'DPM++ SDE Karras'
    DDIM = 'DDIM'
    PLMS = 'PLMS'
    UNIPC = 'UniPC'


# --------------- Refiner ---------------
@dataclass
class Refiner:
    checkpoint: str
    switch_at: float

# --------------- Text2Image ---------------


@dataclass
class Txt2ImgRequest(JSONe):
    prompt: str
    negative_prompt: Optional[str] = None
    model_name: str = 'dreamshaper_5BakedVae.safetensors'
    sampler_name: str = None
    batch_size: int = 1
    n_iter: int = 1
    steps: int = 20
    cfg_scale: float = 7
    height: int = 512
    width: int = 512

    seed: Optional[int] = -1
    restore_faces: Optional[bool] = False
    sd_vae: Optional[str] = None
    clip_skip: Optional[int] = 1

    controlnet_units: Optional[List[ControlnetUnit]] = None
    controlnet_no_detectmap: Optional[bool] = False

    enable_hr: Optional[bool] = False
    hr_upscaler: Optional[str] = 'R-ESRGAN 4x+'
    hr_scale: Optional[float] = 2.0
    hr_resize_x: Optional[int] = None
    hr_resize_y: Optional[int] = None

    sd_refiner: Optional[Refiner] = None


class Txt2ImgResponseCode(Enum):
    NORMAL = 0
    INTERNAL_ERROR = -1
    INVALID_JSON = 1
    MODEL_NOT_EXISTS = 2
    TASK_ID_NOT_EXISTS = 3
    INVALID_AUTH = 4
    HOST_UNAVAILABLE = 5
    PARAM_RANGE_ERROR = 6
    COST_BALANCE_ERROR = 7
    SAMPLER_NOT_EXISTS = 8
    TIMEOUT = 9

    UNKNOWN = 100

    @classmethod
    def _missing_(cls, number):
        return cls(cls.UNKNOWN)


@dataclass
class Txt2ImgResponseData(JSONe):
    task_id: str
    warn: Optional[str] = None


@dataclass
class Txt2ImgResponse(JSONe):
    code: Txt2ImgResponseCode
    msg: str
    data: Optional[Txt2ImgResponseData] = None


#  --------------- Image2Image ---------------


@dataclass
class Img2ImgRequest(JSONe):
    model_name: str = 'dreamshaper_5BakedVae.safetensors'
    sampler_name: str = None
    init_images: List[str] = None
    mask: Optional[str] = None
    resize_mode: Optional[int] = 0
    denoising_strength: Optional[float] = 0.75
    cfg_scale: Optional[float] = None
    mask_blur: Optional[int] = 4
    inpainting_fill: Optional[int] = 1
    inpaint_full_res: Optional[int] = 0
    inpaint_full_res_padding: Optional[int] = 32
    inpainting_mask_invert: Optional[int] = 0
    initial_noise_multiplier: Optional[float] = 1.0
    prompt: Optional[str] = None
    seed: Optional[int] = None
    negative_prompt: Optional[str] = None
    batch_size: Optional[int] = 1
    n_iter: Optional[int] = 1
    steps: Optional[int] = 20
    width: Optional[int] = 1024
    height: Optional[int] = 1024
    restore_faces: Optional[bool] = False
    sd_vae: Optional[str] = None
    clip_skip: Optional[int] = 1

    controlnet_units: Optional[List[ControlnetUnit]] = None
    controlnet_no_detectmap: Optional[bool] = False

    sd_refiner: Optional[Refiner] = None


class Img2ImgResponseCode(Enum):
    NORMAL = 0
    INTERNAL_ERROR = -1
    INVALID_JSON = 1
    MODEL_NOT_EXISTS = 2
    TASK_ID_NOT_EXISTS = 3
    INVALID_AUTH = 4
    HOST_UNAVAILABLE = 5
    PARAM_RANGE_ERROR = 6
    COST_BALANCE_ERROR = 7
    SAMPLER_NOT_EXISTS = 8
    TIMEOUT = 9

    UNKNOWN = 100

    @classmethod
    def _missing_(cls, number):
        return cls(cls.UNKNOWN)


@dataclass
class Img2ImgResponseData(JSONe):
    task_id: str
    warn: Optional[str] = None


@dataclass
class Img2ImgResponse(JSONe):
    code: Img2ImgResponseCode
    msg: str
    data: Optional[Img2ImgResponseData] = None

# --------------- Progress ---------------


class ProgressResponseStatusCode(Enum):
    INITIALIZING = 0
    RUNNING = 1
    SUCCESSFUL = 2
    FAILED = 3
    TIMEOUT = 4

    UNKNOWN = 100

    @classmethod
    def _missing_(cls, number):
        return cls(cls.UNKNOWN)

    def finished(self):
        return self in (ProgressResponseStatusCode.SUCCESSFUL, ProgressResponseStatusCode.FAILED, ProgressResponseStatusCode.TIMEOUT)


@dataclass
class ProgressData(JSONe):
    status: ProgressResponseStatusCode
    progress: int
    eta_relative: int
    imgs: Optional[List[str]] = None
    imgs_bytes: Optional[List[str]] = None
    # info: Optional[Dict[str, str]] = None   # We'll handle this field separately
    failed_reason: Optional[str] = ""
    current_images: Optional[List[str]] = None
    submit_time: Optional[str] = ""
    execution_time: Optional[str] = ""
    txt2img_time: Optional[str] = ""
    finish_time: Optional[str] = ""


class ProgressResponseCode(Enum):
    NORMAL = 0
    INTERNAL_ERROR = -1
    INVALID_JSON = 1
    MODEL_NOT_EXISTS = 2
    TASK_ID_NOT_EXISTS = 3
    INVALID_AUTH = 4
    HOST_UNAVAILABLE = 5
    PARAM_RANGE_ERROR = 6
    COST_BALANCE_ERROR = 7
    SAMPLER_NOT_EXISTS = 8
    TIMEOUT = 9

    UNKNOWN = 100

    @classmethod
    def _missing_(cls, number):
        return cls(cls.UNKNOWN)


@dataclass
class ProgressResponse(JSONe):
    code: ProgressResponseCode
    data: Optional[ProgressData] = None
    msg: Optional[str] = ""

    def download_images(self):
        if self.data.imgs is not None and len(self.data.imgs) > 0:
            self.data.imgs_bytes = batch_download_images(self.data.imgs)

# --------------- Upscale ---------------


class UpscaleResizeMode(Enum):
    SCALE = 0
    SIZE = 1


@dataclass
class UpscaleRequest(JSONe):
    image: str
    upscaler_1: Optional[str] = 'R-ESRGAN 4x+'
    resize_mode: Optional[UpscaleResizeMode] = UpscaleResizeMode.SCALE
    upscaling_resize: Optional[float] = 2.0
    upscaling_resize_w: Optional[int] = None
    upscaling_resize_h: Optional[int] = None
    upscaling_crop: Optional[bool] = False

    upscaler_2: Optional[str] = None
    extras_upscaler_2_visibility: Optional[float] = None
    gfpgan_visibility: Optional[float] = None
    codeformer_visibility: Optional[float] = None
    codeformer_weight: Optional[float] = None


class UpscaleResponseCode(Enum):
    NORMAL = 0
    INTERNAL_ERROR = -1
    INVALID_JSON = 1
    MODEL_NOT_EXISTS = 2
    TASK_ID_NOT_EXISTS = 3
    INVALID_AUTH = 4
    HOST_UNAVAILABLE = 5
    PARAM_RANGE_ERROR = 6
    COST_BALANCE_ERROR = 7
    SAMPLER_NOT_EXISTS = 8
    TIMEOUT = 9

    UNKNOWN = 100

    @classmethod
    def _missing_(cls, number):
        return cls(cls.UNKNOWN)


@dataclass
class UpscaleResponseData(JSONe):
    task_id: str
    warn: Optional[str] = None


@dataclass
class UpscaleResponse(JSONe):
    code: UpscaleResponseCode
    msg: str
    data: Optional[UpscaleResponseData] = None


# --------------- Model ---------------


class ModelType(Enum):
    CHECKPOINT = "checkpoint"
    LORA = "lora"
    VAE = "vae"
    CONTROLNET = "controlnet"
    TEXT_INVERSION = "textualinversion"
    UPSCALER = "upscaler"

    UNKNOWN = "unknown"

    @classmethod
    def _missing_(cls, number):
        return cls(cls.UNKNOWN)


@dataclass
class CivitaiImageMeta(JSONe):
    prompt: Optional[str] = None
    negative_prompt: Optional[str] = None
    sampler_name: Optional[str] = None
    steps: Optional[int] = None
    cfg_scale: Optional[int] = None
    seed: Optional[int] = None
    height: Optional[int] = None
    width: Optional[int] = None
    model_name: Optional[str] = None


@dataclass
class CivitaiImage(JSONe):
    url: str
    nsfw: str
    meta: Optional[CivitaiImageMeta] = None


@dataclass
class ModelInfo(JSONe):
    name: str
    hash: str
    civitai_version_id: int
    sd_name: str
    third_source: str
    download_status: int
    download_name: str
    dependency_status: int
    type: ModelType
    civitai_nsfw: Optional[bool] = False
    civitai_model_id: Optional[int] = 0
    civitai_link: Optional[str] = None
    civitai_images: Optional[List[CivitaiImage]] = field(default_factory=lambda: [])
    civitai_download_url: Optional[str] = None
    civitai_allow_commercial_use: Optional[bool] = True
    civitai_allow_different_license: Optional[bool] = True
    civitai_create_at: Optional[str] = None
    civitai_update_at: Optional[str] = None
    civitai_tags: Optional[str] = None
    civitai_download_count: Optional[int] = 0
    civitai_favorite_count: Optional[int] = 0
    civitai_comment_count: Optional[int] = 0
    civitai_rating_count: Optional[int] = 0
    civitai_rating: Optional[float] = 0.0
    Novita_used_count: Optional[int] = None
    civitai_image_url:  Optional[str] = None
    civitai_image_nsfw:  Optional[bool] = False
    civitai_origin_image_url:  Optional[str] = None
    civitai_image_prompt:  Optional[str] = None
    civitai_image_negative_prompt:  Optional[str] = None
    civitai_image_sampler_name:  Optional[str] = None
    civitai_image_height: Optional[int] = None
    civitai_image_width:  Optional[int] = None
    civitai_image_steps:  Optional[int] = None
    civitai_image_cfg_scale:  Optional[int] = None
    civitai_image_seed:  Optional[int] = None


@dataclass
class ModelData(JSONe):
    models: List[ModelInfo] = None


@dataclass
class MoodelsResponse(JSONe):
    code: int
    msg: str
    data: Optional[ModelData] = field(default_factory=lambda: [])


class ModelList(list):
    """A list of ModelInfo"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def get_by_civitai_version_id(self, civitai_version_id: int):
        for model in self:
            if model.civitai_version_id == civitai_version_id:
                return model
        return None

    def get_by_name(self, name):
        for model in self:
            if model.name == name:
                return model
        return None

    def get_by_sd_name(self, sd_name):
        for model in self:
            if model.sd_name == sd_name:
                return model
        return None

    def list_civitai_tags(self) -> List[str]:
        s = set()
        for model in self:
            if model.civitai_tags:
                s.update(s.strip()
                         for s in model.civitai_tags.split(",") if s.strip())
        return list(s)

    def filter_by_civitai_tags(self, *tags):
        ret = []
        for model in self:
            if model.civitai_tags:
                if set(tags).issubset(set(s.strip() for s in model.civitai_tags.split(","))):
                    ret.append(model)
        return ModelList(ret)

    def filter_by_nsfw(self, nsfw: bool):
        return ModelList([model for model in self if model.civitai_nsfw == nsfw])

    def filter_by_type(self, type):
        return ModelList([model for model in self if model.type == type])

    def filter_by_civitai_model_id(self, civitai_model_id: int):
        return ModelList([model for model in self if model.civitai_model_id == civitai_model_id])

    def filter_by_civitai_model_name(self, name: str):
        return ModelList([model for model in self if model.name == name])

    def sort_by_civitai_download(self):
        return ModelList(sorted(self, key=lambda x: x.civitai_download_count, reverse=True))

    def sort_by_civitai_rating(self):
        return ModelList(sorted(self, key=lambda x: x.civitai_rating, reverse=True))

    def sort_by_civitai_favorite(self):
        return ModelList(sorted(self, key=lambda x: x.civitai_favorite, reverse=True))

    def sort_by_civitai_comment(self):
        return ModelList(sorted(self, key=lambda x: x.civitai_comment, reverse=True))

src/novita_client/proto.py
novitalabs/python-sdk
#!/usr/bin/env python
# -*- coding: UTF-8 -*-

from dataclass_wizard import JSONWizard, DumpMeta


class JSONe(JSONWizard):
    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        DumpMeta(key_transform='SNAKE').bind_to(cls)

src/novita_client/serializer.py
novitalabs/python-sdk
#!/usr/bin/env python
# -*- coding: UTF-8 -*-

DEFAULT_REQUEST_TIMEOUT = 120
DEFAULT_DOWNLOAD_IMAGE_ATTEMPTS = 5
DEFAULT_DOWNLOAD_ONE_IMAGE_TIMEOUT = 30
DEFAULT_POLL_INTERVAL = 0.5
src/novita_client/settings.py
novitalabs/python-sdk
#!/usr/bin/env python
# -*- coding: UTF-8 -*-

from multiprocessing.pool import ThreadPool

import base64
import logging

import requests

from . import settings
from .proto import *

logger = logging.getLogger(__name__)


def batch_download_images(image_links):
    def _download(image_link):
        attempts = settings.DEFAULT_DOWNLOAD_IMAGE_ATTEMPTS
        while attempts > 0:
            try:
                response = requests.get(
                    image_link, timeout=settings.DEFAULT_DOWNLOAD_ONE_IMAGE_TIMEOUT)
                return response.content
            except Exception:
                logger.warning("Failed to download image, retrying...")
            attempts -= 1
        return None

    pool = ThreadPool()
    applied = []
    for img_url in image_links:
        applied.append(pool.apply_async(_download, (img_url, )))
    ret = [r.get() for r in applied]
    return [_ for _ in ret if _ is not None]


def save_image(image_bytes, name):
    with open(name, "wb") as f:
        f.write(image_bytes)


def read_image(name):
    with open(name, "rb") as f:
        return f.read()


def read_image_to_base64(name):
    with open(name, "rb") as f:
        return base64.b64encode(f.read()).decode('utf-8')


def add_lora_to_prompt(prompt: str, lora_name: str, weight: float = 1.0) -> str:
    prompt_split = [s.strip() for s in prompt.split(",")]
    ret = []
    replace = False
    for prompt_chunk in prompt_split:
        if prompt_chunk.startswith("<lora:{}".format(lora_name)):
            ret.append("<lora:{}:{}>".format(lora_name, weight))
            replace = True
        else:
            ret.append(prompt_chunk)
    if not replace:
        ret.append("<lora:{}:{}>".format(lora_name, weight))
    return ", ".join(ret)

src/novita_client/utils.py
novitalabs/python-sdk
from novita_client import *
from novita_client.utils import save_image, read_image_to_base64
import os
from PIL import Image
import io


import pytest


@pytest.mark.dependency()
def test_txt2img_sync():
    client = NovitaClient(os.getenv('NOVITA_API_KEY'))
    res = client.sync_txt2img(Txt2ImgRequest(
        prompt='a dog flying in the sky',
        batch_size=1,
        cfg_scale=7.5,
        sampler_name=Samplers.EULER_A,
    ))

    assert (res.data.status == ProgressResponseStatusCode.SUCCESSFUL)
    assert (len(res.data.imgs_bytes) == 1)

    test_path = os.path.join(os.path.abspath(
        os.path.dirname(__name__)), "tests/data")
    if not os.path.exists(test_path):
        os.makedirs(test_path)
    save_image(res.data.imgs_bytes[0], os.path.join(
        test_path, 'test_txt2img_sync.png'))


@pytest.mark.dependency(depends=['test_txt2img_sync'])
def test_img2img_sync():
    client = NovitaClient(os.getenv('NOVITA_API_KEY'))
    init_image = os.path.join(os.path.abspath(
        os.path.dirname(__name__)), "tests/data/test_txt2img_sync.png")
    init_image_base64 = read_image_to_base64(init_image)

    res = client.sync_img2img(Img2ImgRequest(
        prompt='a dog flying in the sky',
        batch_size=1,
        cfg_scale=7.5,
        sampler_name=Samplers.EULER_A,
        init_images=[init_image_base64]
    ))

    assert (res.data.status == ProgressResponseStatusCode.SUCCESSFUL)
    assert (len(res.data.imgs_bytes) == 1)

    test_path = os.path.join(os.path.abspath(
        os.path.dirname(__name__)), "tests/data")
    if not os.path.exists(test_path):
        os.makedirs(test_path)
    save_image(res.data.imgs_bytes[0], os.path.join(
        test_path, 'test_img2img_sync.png'))


@pytest.mark.dependency(depends=['test_img2img_sync'])
def test_txt2img_controlnet():
    client = NovitaClient(os.getenv('NOVITA_API_KEY'))
    init_image = os.path.join(os.path.abspath(os.path.dirname(
        __name__)), "tests/data/test_txt2img_sync.png")
    init_image_base64 = read_image_to_base64(init_image)

    client = NovitaClient(os.getenv('NOVITA_API_KEY'))
    request = Txt2ImgRequest(
        prompt='a dog flying in the sky',
        batch_size=1,
        cfg_scale=7.5,
        sampler_name=Samplers.EULER_A,
        controlnet_units=[
            ControlnetUnit(
                input_image=init_image_base64,
                model='control_v11p_sd15_canny',
                module='canny',
            ),
        ]
    )

    res = client.sync_txt2img(request)
    test_path = os.path.join(os.path.abspath(
        os.path.dirname(__name__)), "tests/data")

    if not os.path.exists(test_path):
        os.makedirs(test_path)

    assert len(res.data.imgs_bytes) == 2

    save_image(res.data.imgs_bytes[1], os.path.join(
        test_path, 'test_txt2img_controlnet_processor.png'))
    save_image(res.data.imgs_bytes[0], os.path.join(
        test_path, 'test_txt2img_controlnet_result.png'))


@pytest.mark.dependency(depends=['test_img2img_sync'])
def test_img2img_controlnet():
    client = NovitaClient(os.getenv('NOVITA_API_KEY'))
    init_image = os.path.join(os.path.abspath(
        os.path.dirname(__name__)), "tests/data/test_txt2img_sync.png")
    init_image_base64 = read_image_to_base64(init_image)

    res = client.sync_img2img(Img2ImgRequest(
        prompt='a dog flying in the sky',
        batch_size=1,
        cfg_scale=7.5,
        sampler_name=Samplers.EULER_A,
        init_images=[init_image_base64],
        controlnet_units=[
            ControlnetUnit(
                input_image=init_image_base64,
                model='control_v11p_sd15_canny',
                module='canny',
            ),
        ]
    ))

    assert (res.data.status == ProgressResponseStatusCode.SUCCESSFUL)
    assert len(res.data.imgs_bytes) == 2

    test_path = os.path.join(os.path.abspath(
        os.path.dirname(__name__)), "tests/data")
    if not os.path.exists(test_path):
        os.makedirs(test_path)

    save_image(res.data.imgs_bytes[1], os.path.join(
        test_path, 'test_img2img_controlnet_processor.png'))
    save_image(res.data.imgs_bytes[0], os.path.join(
        test_path, 'test_img2img_controlnet_result.png'))


def test_txt2img_upscale_2x():
    client = NovitaClient(os.getenv('NOVITA_API_KEY'))
    res = client.sync_txt2img(Txt2ImgRequest(
        model_name='dreamshaper_8_93211.safetensors',
        prompt='a dog flying in the sky',
        width=512,
        height=512,
        batch_size=1,
        cfg_scale=7.5,
        sampler_name=Samplers.EULER_A,
        enable_hr=True,
        hr_scale=2.0
    ))

    assert (res.data.status == ProgressResponseStatusCode.SUCCESSFUL)
    assert (len(res.data.imgs_bytes) == 1)

    test_path = os.path.join(os.path.abspath(
        os.path.dirname(__name__)), "tests/data")
    if not os.path.exists(test_path):
        os.makedirs(test_path)
    img = Image.open(io.BytesIO(res.data.imgs_bytes[0]))
    assert img.size == (1024, 1024)


def test_txt2img_upscale_specify_size():
    client = NovitaClient(os.getenv('NOVITA_API_KEY'))
    res = client.sync_txt2img(Txt2ImgRequest(
        model_name='dreamshaper_8_93211.safetensors',
        prompt='a dog flying in the sky',
        width=512,
        height=512,
        batch_size=1,
        cfg_scale=7.5,
        sampler_name=Samplers.EULER_A,
        enable_hr=True,
        hr_resize_x=768,
        hr_resize_y=768
    ))

    assert (res.data.status == ProgressResponseStatusCode.SUCCESSFUL)
    assert (len(res.data.imgs_bytes) == 1)

    test_path = os.path.join(os.path.abspath(
        os.path.dirname(__name__)), "tests/data")
    if not os.path.exists(test_path):
        os.makedirs(test_path)
    img = Image.open(io.BytesIO(res.data.imgs_bytes[0]))
    assert img.size == (768, 768)


def test_upscale_2x():
    client = NovitaClient(os.getenv('NOVITA_API_KEY'))
    res = client.sync_img2img(Img2ImgRequest(
        model_name='dreamshaper_8_93211.safetensors',
        prompt='a dog flying in the sky',
        width=512,
        height=512,
        batch_size=1,
        cfg_scale=7.5,
        sampler_name=Samplers.EULER_A,
    ))

    assert (res.data.status == ProgressResponseStatusCode.SUCCESSFUL)
    assert (len(res.data.imgs_bytes) == 1)

    image = base64.b64encode(res.data.imgs_bytes[0]).decode('utf-8')
    upscale_req = UpscaleRequest(
        image=image,
        resize_mode=UpscaleResizeMode.SCALE,
        upscaling_resize=2
    )
    upscale_res = client.sync_upscale(upscale_req)
    assert (upscale_res.data.status == ProgressResponseStatusCode.SUCCESSFUL)
    assert (len(upscale_res.data.imgs_bytes) == 1)
    img = Image.open(io.BytesIO(upscale_res.data.imgs_bytes[0]))
    assert img.size == (1024, 1024)


def test_upscale_specify_size():
    client = NovitaClient(os.getenv('NOVITA_API_KEY'))
    res = client.sync_img2img(Img2ImgRequest(
        model_name='dreamshaper_8_93211.safetensors',
        prompt='a dog flying in the sky',
        width=512,
        height=512,
        batch_size=1,
        cfg_scale=7.5,
        sampler_name=Samplers.EULER_A,
    ))

    assert (res.data.status == ProgressResponseStatusCode.SUCCESSFUL)
    assert (len(res.data.imgs_bytes) == 1)

    image = base64.b64encode(res.data.imgs_bytes[0]).decode('utf-8')
    upscale_req = UpscaleRequest(
        image=image,
        resize_mode=UpscaleResizeMode.SIZE,
        upscaling_resize_h=768,
        upscaling_resize_w=768
    )
    upscale_res = client.sync_upscale(upscale_req)
    assert (upscale_res.data.status == ProgressResponseStatusCode.SUCCESSFUL)
    assert (len(upscale_res.data.imgs_bytes) == 1)
    img = Image.open(io.BytesIO(upscale_res.data.imgs_bytes[0]))
    assert img.size == (768, 768)


def test_upscale_multiple_upscaler():
    client = NovitaClient(os.getenv('NOVITA_API_KEY'))
    res = client.sync_img2img(Img2ImgRequest(
        model_name='dreamshaper_8_93211.safetensors',
        prompt='a dog flying in the sky',
        width=512,
        height=512,
        batch_size=1,
        cfg_scale=7.5,
        sampler_name=Samplers.EULER_A,
    ))

    assert (res.data.status == ProgressResponseStatusCode.SUCCESSFUL)
    assert (len(res.data.imgs_bytes) == 1)

    image = base64.b64encode(res.data.imgs_bytes[0]).decode('utf-8')
    upscale_req = UpscaleRequest(
        image=image,
        upscaler_1='R-ESRGAN 4x+',
        resize_mode=UpscaleResizeMode.SIZE,
        upscaling_resize_h=768,
        upscaling_resize_w=768,
        upscaler_2='Nearest',
        gfpgan_visibility=1,
    )
    upscale_res = client.sync_upscale(upscale_req)
    assert (upscale_res.data.status == ProgressResponseStatusCode.SUCCESSFUL)
    assert (len(upscale_res.data.imgs_bytes) == 1)
    img = Image.open(io.BytesIO(upscale_res.data.imgs_bytes[0]))
    assert img.size == (768, 768)


def test_txt2img_custom_headers():
    client = NovitaClient(os.getenv('NOVITA_API_KEY'))
    client.set_extra_headers({"User-Agent": "test-custom-user-agent"})

    res = client.sync_img2img(Img2ImgRequest(
        model_name='dreamshaper_8_93211.safetensors',
        prompt='a dog flying in the sky',
        width=512,
        height=512,
        batch_size=1,
        cfg_scale=7.5,
        sampler_name=Samplers.EULER_A,
    ))

    assert (res.data.status == ProgressResponseStatusCode.SUCCESSFUL)
    assert (len(res.data.imgs_bytes) == 1)


def test_txt2img_with_callback():
    client = NovitaClient(os.getenv('NOVITA_API_KEY'))

    def callback(res: ProgressResponse):
        assert isinstance(res.data.progress, float)

    res = client.sync_txt2img(Txt2ImgRequest(
        model_name='dreamshaper_8_93211.safetensors',
        prompt='a dog flying in the sky',
        width=512,
        height=512,
        batch_size=1,
        cfg_scale=7.5,
        sampler_name=Samplers.EULER_A,
    ), callback=callback)

    assert (res.data.status == ProgressResponseStatusCode.SUCCESSFUL)
    assert (len(res.data.imgs_bytes) == 1)


def test_txt2img_sdxl_with_refiner():
    client = NovitaClient(os.getenv('NOVITA_API_KEY'))
    res = client.sync_txt2img(Txt2ImgRequest(
        model_name='sd_xl_base_1.0.safetensors',
        prompt='a dog flying in the sky',
        width=1024,
        height=1024,
        batch_size=1,
        cfg_scale=7.5,
        sampler_name=Samplers.EULER_A,
        sd_refiner=Refiner(
            checkpoint='sd_xl_refiner_1.0.safetensors',
            switch_at=0.5,
        )
    ))

    assert (res.data.status == ProgressResponseStatusCode.SUCCESSFUL)
    assert (len(res.data.imgs_bytes) == 1)

    test_path = os.path.join(os.path.abspath(
        os.path.dirname(__name__)), "tests/data")
    if not os.path.exists(test_path):
        os.makedirs(test_path)
    save_image(res.data.imgs_bytes[0], os.path.join(
        test_path, 'test_sdxl_txt2img_refienr.png'))


def test_img2img_sdxl_with_refiner():
    client = NovitaClient(os.getenv('NOVITA_API_KEY'))
    res = client.sync_txt2img(Txt2ImgRequest(
        model_name='sd_xl_base_1.0.safetensors',
        prompt='a dog flying in the sky',
        width=1024,
        height=1024,
        batch_size=1,
        cfg_scale=7.5,
        sampler_name=Samplers.EULER_A,
        sd_refiner=Refiner(
            checkpoint='sd_xl_refiner_1.0.safetensors',
            switch_at=0.5,
        )
    ))

    assert (res.data.status == ProgressResponseStatusCode.SUCCESSFUL)
    assert (len(res.data.imgs_bytes) == 1)

    base64_image = base64.b64encode(res.data.imgs_bytes[0]).decode('utf-8')

    res = client.sync_img2img(Img2ImgRequest(
        prompt='a dog flying in the sky',
        model_name='sd_xl_base_1.0.safetensors',
        batch_size=1,
        cfg_scale=7.5,
        sampler_name=Samplers.EULER_A,
        init_images=[base64_image],
        sd_refiner=Refiner(
            checkpoint='sd_xl_refiner_1.0.safetensors',
            switch_at=0.5,
        )
    ))

    assert (res.data.status == ProgressResponseStatusCode.SUCCESSFUL)
    assert (len(res.data.imgs_bytes) == 1)

    test_path = os.path.join(os.path.abspath(
        os.path.dirname(__name__)), "tests/data")
    if not os.path.exists(test_path):
        os.makedirs(test_path)
    save_image(res.data.imgs_bytes[0], os.path.join(
        test_path, 'test_sdxl_img2img_refienr.png'))

tests/test_basics.py
novitalabs/python-sdk
from novita_client import *
import os


def test_model_api():
    client = NovitaClient(os.getenv('NOVITA_API_KEY'))
    models = client.models()
    assert all([m.civitai_nsfw is True for m in models.filter_by_nsfw(True)])
    assert all([m.civitai_nsfw is False for m in models.filter_by_nsfw(False)])

    assert len(models. \
        filter_by_type(ModelType.LORA). \
        filter_by_nsfw(False). \
        filter_by_civitai_tags('anime'). \
        sort_by_civitai_rating()) > 0

    assert len(models.filter_by_type(ModelType.CHECKPOINT)) > 0
    assert len(models.filter_by_type(ModelType.LORA)) > 0
    assert len(models.filter_by_type(ModelType.TEXT_INVERSION)) > 0
    assert len(models.filter_by_type(ModelType.CONTROLNET)) > 0

tests/test_model.py
novitalabs/python-sdk