Merge pull request #50 from AIrjen/api_call_robustness

Added API stability

call_extras.py

@@ -4,6 +4,7 @@ import io
 import base64
 import uuid
 from PIL import Image, PngImagePlugin
+import time
 
 
 
@@ -46,8 +47,25 @@ def call_extras(imagelocation,originalimage, originalpnginfo ="", apiurl="http:/
         "models": "None" # the remove backgrounds plugin is  automatically turned on, need to turn it off
     }
 
+    response = []
+        
+    # If we don't get an image back, we want to retry a few times. Max 3 times
+    for i in range(4):
+        response = requests.post(url=f'{url}/sdapi/v1/extra-single-image', json=payload)
 
-    response = requests.post(url=f'{url}/sdapi/v1/extra-single-image', json=payload)
+        r = response.json()
+        if("image" in r):
+            break # this means if we have the images object, then we "break" out of the for loop.
+        else:
+            if(i == 3):
+                print("If this keeps happening: Is WebUI started with --api enabled?")
+                print("")
+                raise ValueError("API has not been responding after several retries. Stopped processing.")
+            print("")
+            print("We haven't received an image from the API. Maybe something went wrong. Will retry after waiting a bit.")
+            
+
+            time.sleep(10 * (i+1) ) # incremental waiting time
 
     image = Image.open(io.BytesIO(base64.b64decode(response.json().get("image"))))
 

call_img2img.py

@@ -6,6 +6,7 @@ import uuid
 from PIL import Image, PngImagePlugin
 from modules import shared
 from model_lists import *
+import time
 
 
 def call_img2img(imagelocation,originalimage, originalpnginfo ="", apiurl="http://127.0.0.1:7860",filename="", prompt = "", negativeprompt = "", img2imgsamplingsteps = "20", img2imgcfg = "7", img2imgsamplingmethod = "DPM++ SDE Karras", img2imgupscaler = "R-ESRGAN 4x+", img2imgmodel = "currently selected model", denoising_strength = "0.3", scale = "2", padding = "64",upscalescript="SD upscale",usdutilewidth = "512", usdutileheight = "0", usdumaskblur = "8", usduredraw ="Linear", usduSeamsfix = "None", usdusdenoise = "0.35", usduswidth = "64", usduspadding ="32", usdusmaskblur = "8",controlnetenabled=False, controlnetmodel="",controlnetblockymode=False):
@@ -171,13 +172,25 @@ def call_img2img(imagelocation,originalimage, originalpnginfo ="", apiurl="http:
     # target_size_type = 2
     # custom_scale = 2
 
+    r = []
+    # If we don't get an image back, we want to retry a few times. Max 3 times
+    for i in range(4):
+        response = requests.post(url=f'{url}/sdapi/v1/img2img', json=payload)
 
 
-    response = requests.post(url=f'{url}/sdapi/v1/img2img', json=payload)
-
-
-    r = response.json()
+        r = response.json()
+        if('images' in r):
+            break # this means if we have the images object, then we "break" out of the for loop.
+        else:
+            if(i == 3):
+                print("If this keeps happening: Is WebUI started with --api enabled?")
+                print("")
+                raise ValueError("API has not been responding after several retries. Stopped processing.")
+            print("")
+            print("We haven't received an image from the API. Maybe something went wrong. Will retry after waiting a bit.")
+                
 
+            time.sleep(10 * (i+1) ) # incremental waiting time
 
     for i in r['images']:
         image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0])))

call_txt2img.py

@@ -6,6 +6,7 @@ import uuid
 import sys, os
 from PIL import Image, PngImagePlugin
 from model_lists import *
+import time
 
 def call_txt2img(passingprompt,ratio,upscale,debugmode,filename="",model = "currently selected model",samplingsteps = "40",cfg= "7",hiressteps ="0",denoisestrength="0.6",samplingmethod="DPM++ SDE Karras", upscaler="R-ESRGAN 4x+",hiresscale="2",apiurl="http://127.0.0.1:7860", qualitygate=False,quality="7.6",runs="5",negativeprompt=""):
 
@@ -130,10 +131,28 @@ def call_txt2img(passingprompt,ratio,upscale,debugmode,filename="",model = "curr
         filenamefull = filename + addrun
         outputTXT2IMGFull = '{}{}{}'.format(outputTXT2IMGfolder,filenamefull,outputTXT2IMGpng)
 
+        r = []
+        
+        # If we don't get an image back, we want to retry a few times. Max 3 times
+        for i in range(4):
+            response = requests.post(url=f'{url}/sdapi/v1/txt2img', json=payload)
+
+            r = response.json()
+            if('images' in r):
+                break # this means if we have the images object, then we "break" out of the for loop.
+            else:
+                if(i == 3):
+                    print("If this keeps happening: Is WebUI started with --api enabled?")
+                    print("")
+                    raise ValueError("API has not been responding after several retries. Stopped processing.")
+                print("")
+                print("We haven't received an image from the API. Maybe something went wrong. Will retry after waiting a bit.")
+                
+
+                time.sleep(10 * (i+1) ) # incremental waiting time
+
 
-        response = requests.post(url=f'{url}/sdapi/v1/txt2img', json=payload)
 
-        r = response.json()
 
         for i in r['images']:
             image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0])))

model_lists.py

@@ -25,6 +25,7 @@ def get_samplers():
 
 def get_upscalers_for_img2img():
     upscalerlistfromwWebUI = [x.name for x in shared.sd_upscalers]
+    upscalerlistfromwWebUI.remove("None")
     return upscalerlistfromwWebUI
 
 def get_samplers_for_img2img():