Add "extra" single image operation

import io

import base64

from modules.api.models import *

from PIL import Image

from modules.extras import run_extras

def upscaler_to_index(name: str):

    try:

        return [x.name.lower() for x in shared.sd_upscalers].index(name.lower())

    except:

        raise HTTPException(status_code=400, detail="Upscaler not found")

sampler_to_index = lambda name: next(filter(lambda row: name.lower() == row[1].name.lower(), enumerate(all_samplers)), None)

def img_to_base64(img):

def img_to_base64(img: str):

    buffer = io.BytesIO()

    img.save(buffer, format="png")

    return base64.b64encode(buffer.getvalue())

def base64_to_bytes(base64Img: str):

    if "," in base64Img:

        base64Img = base64Img.split(",")[1]

    return io.BytesIO(base64.b64decode(base64Img))

def base64_to_images(base64Imgs: list[str]):

    imgs = []

    for img in base64Imgs:

        img = Image.open(base64_to_bytes(img))

        imgs.append(img)

    return imgs

class Api:

    def __init__(self, app, queue_lock):

        self.router = APIRouter()

        self.app = app

        self.queue_lock = queue_lock

        self.app.add_api_route("/sdapi/v1/txt2img", self.text2imgapi, methods=["POST"])

        self.app.add_api_route("/sdapi/v1/txt2img", self.text2imgapi, methods=["POST"], response_model=TextToImageResponse)

        self.app.add_api_route("/sdapi/v1/extra-single-image", self.extras_single_image_api, methods=["POST"], response_model=ExtrasSingleImageResponse)

    def text2imgapi(self, txt2imgreq: StableDiffusionProcessingAPI ):

        sampler_index = sampler_to_index(txt2imgreq.sampler_index)

        return TextToImageResponse(images=b64images, parameters=json.dumps(vars(txt2imgreq)), info=json.dumps(processed.info))

    def img2imgapi(self):

        raise NotImplementedError

    def extrasapi(self):

        raise NotImplementedError

    def extras_single_image_api(self, req: ExtrasSingleImageRequest):

        upscaler1Index = upscaler_to_index(req.upscaler_1)

        upscaler2Index = upscaler_to_index(req.upscaler_2)

        reqDict = vars(req)

        reqDict.pop('upscaler_1')

        reqDict.pop('upscaler_2')

        reqDict['image'] = base64_to_images([reqDict['image']])[0]

        with self.queue_lock:

            result = run_extras(**reqDict, extras_upscaler_1=upscaler1Index, extras_upscaler_2=upscaler2Index, extras_mode=0, image_folder="", input_dir="", output_dir="")

        return ExtrasSingleImageResponse(image="data:image/png;base64,"+img_to_base64(result[0]), html_info_x=result[1], html_info=result[2])

    def pnginfoapi(self):

        raise NotImplementedError

from pydantic import BaseModel, Field, Json

from typing_extensions import Literal

from modules.shared import sd_upscalers

class TextToImageResponse(BaseModel):

    images: list[str] = Field(default=None, title="Image", description="The generated image in base64 format.")

    parameters: Json

    info: Json

class ExtrasBaseRequest(BaseModel):

    resize_mode: Literal[0, 1] = Field(default=0, title="Resize Mode", description="Sets the resize mode: 0 to upscale by upscaling_resize amount, 1 to upscale up to upscaling_resize_h x upscaling_resize_w.")

    show_extras_results: bool = Field(default=True, title="Show results", description="Should the backend return the generated image?")

    gfpgan_visibility: float = Field(default=0, title="GFPGAN Visibility", ge=0, le=1, allow_inf_nan=False, description="Sets the visibility of GFPGAN, values should be between 0 and 1.")

    codeformer_visibility: float = Field(default=0, title="CodeFormer Visibility", ge=0, le=1, allow_inf_nan=False, description="Sets the visibility of CodeFormer, values should be between 0 and 1.")

    codeformer_weight: float = Field(default=0, title="CodeFormer Weight", ge=0, le=1, allow_inf_nan=False, description="Sets the weight of CodeFormer, values should be between 0 and 1.")

    upscaling_resize: float = Field(default=2, title="Upscaling Factor", ge=1, le=4, description="By how much to upscale the image, only used when resize_mode=0.")

    upscaling_resize_w: int = Field(default=512, title="Target Width", ge=1, description="Target width for the upscaler to hit. Only used when resize_mode=1.")

    upscaling_resize_h: int = Field(default=512, title="Target Height", ge=1, description="Target height for the upscaler to hit. Only used when resize_mode=1.")

    upscaling_crop: bool = Field(default=True, title="Crop to fit", description="Should the upscaler crop the image to fit in the choosen size?")

    upscaler_1: str = Field(default="None", title="Main upscaler", description=f"The name of the main upscaler to use, it has to be one of this list: {' , '.join([x.name for x in sd_upscalers])}")

    upscaler_2: str = Field(default="None", title="Secondary upscaler", description=f"The name of the secondary upscaler to use, it has to be one of this list: {' , '.join([x.name for x in sd_upscalers])}")

    extras_upscaler_2_visibility: float = Field(default=0, title="Secondary upscaler visibility", ge=0, le=1, allow_inf_nan=False, description="Sets the visibility of secondary upscaler, values should be between 0 and 1.")

class ExtraBaseResponse(BaseModel):

    html_info_x: str

    html_info: str

class ExtrasSingleImageRequest(ExtrasBaseRequest):

    image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.")

class ExtrasSingleImageResponse(ExtraBaseResponse):

    image: str = Field(default=None, title="Image", description="The generated image in base64 format.")

 No newline at end of file