From 97f7e79391c27474b3c2096d4c59ab6e664b27a5 Mon Sep 17 00:00:00 2001
From: AyaanShah2204 <89650087+AyaanShah2204@users.noreply.github.com>
Date: Thu, 22 Jun 2023 16:08:26 -0700
Subject: [PATCH] [Blender Integration] Fixed Inpainting REST API (#1577)

* fixed inpaint api

* added inpainting test

* fixed linter errors

---------

Co-authored-by: powderluv <powderluv@users.noreply.github.com>
---
 apps/stable_diffusion/web/ui/inpaint_ui.py |  6 +-
 rest_api_tests/api_test.py                 | 69 +++++++++++++++++++---
 2 files changed, 65 insertions(+), 10 deletions(-)

diff --git a/apps/stable_diffusion/web/ui/inpaint_ui.py b/apps/stable_diffusion/web/ui/inpaint_ui.py
index 082668e85b..be8a58def4 100644
--- a/apps/stable_diffusion/web/ui/inpaint_ui.py
+++ b/apps/stable_diffusion/web/ui/inpaint_ui.py
@@ -278,7 +278,7 @@ def inpaint_api(
         custom_model="None",
         hf_model_id=InputData["hf_model_id"]
         if "hf_model_id" in InputData.keys()
-        else "stabilityai/stable-diffusion-2-1-base",
+        else "stabilityai/stable-diffusion-2-inpainting",
         custom_vae="None",
         precision="fp16",
         device=available_devices[0],
@@ -289,6 +289,10 @@ def inpaint_api(
         lora_hf_id="",
         ondemand=False,
     )
+
+    # Converts generator type to subscriptable
+    res = next(res)
+
     return {
         "images": encode_pil_to_base64(res[0]),
         "parameters": {},
diff --git a/rest_api_tests/api_test.py b/rest_api_tests/api_test.py
index 844dcfd58b..2265e06d26 100644
--- a/rest_api_tests/api_test.py
+++ b/rest_api_tests/api_test.py
@@ -91,19 +91,70 @@ def img2img_test():
 
     print("Extracting response object")
 
-    # Uncomment below to save the picture
+    # NOTE Uncomment below to save the picture
 
-    response_obj = res.json()
-    img_b64 = response_obj.get("images", [False])[0] or response_obj.get(
-        "image"
+    # response_obj = res.json()
+    # img_b64 = response_obj.get("images", [False])[0] or response_obj.get(
+    #     "image"
+    # )
+    # img_b2 = base64.b64decode(img_b64.replace("data:image/png;base64,", ""))
+    # im_file = BytesIO(img_b2)
+    # response_img = Image.open(im_file)
+    # print("Saving Response Image to: response_img")
+    # response_img.save(r"rest_api_tests/response_img.png")
+
+
+def inpainting_test():
+    prompt = "Paint a rabbit riding on the dog"
+    negative_prompt = "ugly, bad art, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, blurry, bad anatomy, blurred, watermark, grainy, tiling, signature, cut off, draft"
+    seed = 2121991605
+    height = 512
+    width = 512
+    steps = 50
+    noise_level = 10
+    cfg_scale = 7
+    is_full_res = False
+    full_res_padding = 32
+    image_path = r"./rest_api_tests/dog.png"
+
+    img_file = open(image_path, "rb")
+    image = (
+        "data:image/png;base64," + base64.b64encode(img_file.read()).decode()
     )
-    img_b2 = base64.b64decode(img_b64.replace("data:image/png;base64,", ""))
-    im_file = BytesIO(img_b2)
-    response_img = Image.open(im_file)
-    print("Saving Response Image to: response_img")
-    response_img.save(r"rest_api_tests/response_img.png")
+    img_file = open(image_path, "rb")
+    mask = (
+        "data:image/png;base64," + base64.b64encode(img_file.read()).decode()
+    )
+
+    url = "http://127.0.0.1:8080/sdapi/v1/inpaint"
+
+    headers = {
+        "User-Agent": "PythonTest",
+        "Accept": "*/*",
+        "Accept-Encoding": "gzip, deflate, br",
+    }
+
+    data = {
+        "prompt": prompt,
+        "negative_prompt": negative_prompt,
+        "image": image,
+        "mask": mask,
+        "height": height,
+        "width": width,
+        "steps": steps,
+        "noise_level": noise_level,
+        "cfg_scale": cfg_scale,
+        "seed": seed,
+        "is_full_res": is_full_res,
+        "full_res_padding": full_res_padding,
+    }
+
+    res = requests.post(url=url, json=data, headers=headers, timeout=1000)
+
+    print(f"[Inpainting] response from server was : {res.status_code}")
 
 
 if __name__ == "__main__":
     img2img_test()
     upscaler_test()
+    inpainting_test()