From 4ff852ffb50859f2eae75375cab94dd790a46886 Mon Sep 17 00:00:00 2001
From: Bruno Seoane <brunoseoaneamarillo@gmail.com>
Date: Sun, 23 Oct 2022 13:07:59 -0300
Subject: [PATCH] Add batch processing "extras" endpoint

---
 modules/api/api.py    | 25 +++++++++++++++++++++++--
 modules/api/models.py | 15 ++++++++++++++-
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/modules/api/api.py b/modules/api/api.py
index 3b804373..528134a8 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -10,6 +10,7 @@ import base64
 from modules.api.models import *
 from PIL import Image
 from modules.extras import run_extras
+from gradio import processing_utils
 
 def upscaler_to_index(name: str):
     try:
@@ -44,6 +45,7 @@ class Api:
         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/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 text2imgapi(self, txt2imgreq: StableDiffusionProcessingAPI ):
         sampler_index = sampler_to_index(txt2imgreq.sampler_index)
@@ -78,12 +80,31 @@ class Api:
         reqDict.pop('upscaler_1')
         reqDict.pop('upscaler_2')
 
-        reqDict['image'] = base64_to_images([reqDict['image']])[0]
+        reqDict['image'] = processing_utils.decode_base64_to_file(reqDict['image'])
 
         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])
+        return ExtrasSingleImageResponse(image=processing_utils.encode_pil_to_base64(result[0]), html_info_x=result[1], html_info=result[2])
+
+    def extras_batch_images_api(self, req: ExtrasBatchImagesRequest):
+        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_folder'] = list(map(processing_utils.decode_base64_to_file, reqDict['imageList']))
+        reqDict.pop('imageList')
+
+        with self.queue_lock:
+            result = run_extras(**reqDict, extras_upscaler_1=upscaler1Index, extras_upscaler_2=upscaler2Index, extras_mode=1, image="", input_dir="", output_dir="")
+
+        return ExtrasBatchImagesResponse(images=list(map(processing_utils.encode_pil_to_base64, result[0])), html_info_x=result[1], html_info=result[2])
+    
+    def extras_folder_processing_api(self):
+        raise NotImplementedError
 
     def pnginfoapi(self):
         raise NotImplementedError
diff --git a/modules/api/models.py b/modules/api/models.py
index dcf1ab54..bbd0ef53 100644
--- a/modules/api/models.py
+++ b/modules/api/models.py
@@ -29,4 +29,17 @@ 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
+    image: str = Field(default=None, title="Image", description="The generated image in base64 format.")
+
+class SerializableImage(BaseModel):
+  path: str = Field(title="Path", description="The image's path ()")
+
+class ImageItem(BaseModel):
+    data: str = Field(title="image data")
+    name: str = Field(title="filename")
+
+class ExtrasBatchImagesRequest(ExtrasBaseRequest):
+    imageList: list[str] = Field(title="Images", description="List of images to work on. Must be Base64 strings")
+
+class ExtrasBatchImagesResponse(ExtraBaseResponse):
+    images: list[str] = Field(title="Images", description="The generated images in base64 format.")
\ No newline at end of file