mrq/stable-diffusion-webui
1
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update endpoints to use gradio's own utils functions

Browse Source
This commit is contained in:
Bruno Seoane 2022-10-23 15:13:37 -03:00
parent e3f0e34cd6
commit e0ca4dfbc1
2 changed files with 36 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

71
modules/api/api.py
View File

@ -20,27 +20,27 @@ def upscaler_to_index(name: str):
sampler_to_index = lambda name: next(filter(lambda row: name.lower() == row[1].name.lower(), enumerate(all_samplers)), None)
def img_to_base64(img: str):
buffer = io.BytesIO()
img.save(buffer, format="png")
return base64.b64encode(buffer.getvalue())
# 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_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
# 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 ImageToImageResponse(BaseModel):
images: list[str] = Field(default=None, title="Image", description="The generated image in base64 format.")
parameters: Json
info: Json
parameters: dict
info: str
class Api:
@ -49,17 +49,17 @@ class Api:
self.app = app
self.queue_lock = queue_lock
self.app.add_api_route("/sdapi/v1/txt2img", self.text2imgapi, methods=["POST"], response_model=TextToImageResponse)
self.app.add_api_route("/sdapi/v1/img2img", self.img2imgapi, methods=["POST"])
self.app.add_api_route("/sdapi/v1/img2img", self.img2imgapi, methods=["POST"], response_model=ImageToImageResponse)
self.app.add_api_route("/sdapi/v1/extra-single-image", self.extras_single_image_api, methods=["POST"], response_model=ExtrasSingleImageResponse)
self.app.add_api_route("/sdapi/v1/extra-batch-image", self.extras_batch_images_api, methods=["POST"], response_model=ExtrasBatchImagesResponse)
def __base64_to_image(self, base64_string):
# if has a comma, deal with prefix
if "," in base64_string:
base64_string = base64_string.split(",")[1]
imgdata = base64.b64decode(base64_string)
# convert base64 to PIL image
return Image.open(io.BytesIO(imgdata))
# def __base64_to_image(self, base64_string):
# # if has a comma, deal with prefix
# if "," in base64_string:
# base64_string = base64_string.split(",")[1]
# imgdata = base64.b64decode(base64_string)
# # convert base64 to PIL image
# return Image.open(io.BytesIO(imgdata))
def text2imgapi(self, txt2imgreq: StableDiffusionTxt2ImgProcessingAPI):
sampler_index = sampler_to_index(txt2imgreq.sampler_index)
@ -79,11 +79,9 @@ class Api:
with self.queue_lock:
processed = process_images(p)
b64images = list(map(img_to_base64, processed.images))
return TextToImageResponse(images=b64images, parameters=json.dumps(vars(txt2imgreq)), info=json.dumps(processed.info))
b64images = list(map(processing_utils.encode_pil_to_base64, processed.images))
return TextToImageResponse(images=b64images, parameters=json.dumps(vars(txt2imgreq)), info=processed.info)
def img2imgapi(self, img2imgreq: StableDiffusionImg2ImgProcessingAPI):
sampler_index = sampler_to_index(img2imgreq.sampler_index)
@ -98,7 +96,7 @@ class Api:
mask = img2imgreq.mask
if mask:
mask = self.__base64_to_image(mask)
mask = processing_utils.decode_base64_to_image(mask)
populate = img2imgreq.copy(update={ # Override __init__ params
@ -113,7 +111,7 @@ class Api:
imgs = []
for img in init_images:
img = self.__base64_to_image(img)
img = processing_utils.decode_base64_to_image(img)
imgs = [img] * p.batch_size
p.init_images = imgs
@ -121,13 +119,12 @@ class Api:
with self.queue_lock:
processed = process_images(p)
b64images = []
for i in processed.images:
buffer = io.BytesIO()
i.save(buffer, format="png")
b64images.append(base64.b64encode(buffer.getvalue()))
return ImageToImageResponse(images=b64images, parameters=json.dumps(vars(img2imgreq)), info=json.dumps(processed.info))
b64images = list(map(processing_utils.encode_pil_to_base64, processed.images))
# for i in processed.images:
# buffer = io.BytesIO()
# i.save(buffer, format="png")
# b64images.append(base64.b64encode(buffer.getvalue()))
return ImageToImageResponse(images=b64images, parameters=vars(img2imgreq), info=processed.info)
def extras_single_image_api(self, req: ExtrasSingleImageRequest):
upscaler1Index = upscaler_to_index(req.upscaler_1)

4
modules/api/models.py
View File

@ -4,8 +4,8 @@ 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
parameters: str
info: str
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.")