add support to run it on sagemaker

Author: xieyongliang

modules/api/api.py

@@ -77,6 +77,8 @@ def __init__(self, app: FastAPI, queue_lock: Lock):
         self.app.add_api_route("/sdapi/v1/prompt-styles", self.get_promp_styles, methods=["GET"], response_model=List[PromptStyleItem])
         self.app.add_api_route("/sdapi/v1/artist-categories", self.get_artists_categories, methods=["GET"], response_model=List[str])
         self.app.add_api_route("/sdapi/v1/artists", self.get_artists, methods=["GET"], response_model=List[ArtistItem])
+        self.app.add_api_route("/invocations", self.invocations, methods=["POST"], response_model=InvocationsResponse)
+        self.app.add_api_route("/ping", self.ping, methods=["GET"], response_model=PingResponse)
 
     def text2imgapi(self, txt2imgreq: StableDiffusionTxt2ImgProcessingAPI):
         sampler_index = sampler_to_index(txt2imgreq.sampler_index)
@@ -92,13 +94,17 @@ def text2imgapi(self, txt2imgreq: StableDiffusionTxt2ImgProcessingAPI):
             }
         )
         p = StableDiffusionProcessingTxt2Img(**vars(populate))
+
         # Override object param
 
         shared.state.begin()
 
         with self.queue_lock:
-            processed = process_images(p)
-
+            if p.script_args is not None:
+                processed = p.scripts.run(p, *p.script_args)
+            if processed is None:
+                processed = process_images(p)
+        
         shared.state.end()
 
         b64images = list(map(encode_pil_to_base64, processed.images))
@@ -141,7 +147,10 @@ def img2imgapi(self, img2imgreq: StableDiffusionImg2ImgProcessingAPI):
         shared.state.begin()
 
         with self.queue_lock:
-            processed = process_images(p)
+            if p.script_args is not None:
+                processed = p.scripts.run(p, *p.script_args)
+            if processed is None:
+                processed = process_images(p)
 
         shared.state.end()
 
@@ -297,6 +306,17 @@ def get_artists_categories(self):
     def get_artists(self):
         return [{"name":x[0], "score":x[1], "category":x[2]} for x in shared.artist_db.artists]
 
+    def invocations(self, req: InvocationsRequest):
+        if req.task == 'text-to-image':
+            return self.text2imgapi(req.payload)
+        elif req.task == 'image-to-image':
+            return self.img2imgapi(req.payload)
+        else:
+            raise NotImplementedError
+    
+    def ping(self):
+        return {'status': 'Healthy'}
+
     def launch(self, server_name, port):
         self.app.include_router(self.router)
-        uvicorn.run(self.app, host=server_name, port=port)
+        uvicorn.run(self.app, host=server_name, port=port)

modules/api/models.py

@@ -6,6 +6,7 @@
 from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img
 from modules.shared import sd_upscalers, opts, parser
 from typing import Dict, List
+from typing import Union
 
 API_NOT_ALLOWED = [
     "self",
@@ -239,3 +240,14 @@ class ArtistItem(BaseModel):
     score: float = Field(title="Score")
     category: str = Field(title="Category")
 
+class InvocationsRequest(BaseModel):
+    task: str
+    payload: Union[StableDiffusionTxt2ImgProcessingAPI, StableDiffusionImg2ImgProcessingAPI]
+
+class InvocationsResponse(BaseModel):
+    images: list[str] = Field(default=None, title="Image", description="The generated image in base64 format.")
+    parameters: dict
+    info: str
+
+class PingResponse(BaseModel):
+    status: str

modules/processing.py

@@ -78,10 +78,10 @@ class StableDiffusionProcessing():
     """
     The first set of paramaters: sd_models -> do_not_reload_embeddings represent the minimum required to create a StableDiffusionProcessing
     """
-    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt: str = "", styles: List[str] = None, seed: int = -1, subseed: int = -1, subseed_strength: float = 0, seed_resize_from_h: int = -1, seed_resize_from_w: int = -1, seed_enable_extras: bool = True, sampler_index: int = 0, batch_size: int = 1, n_iter: int = 1, steps: int = 50, cfg_scale: float = 7.0, width: int = 512, height: int = 512, restore_faces: bool = False, tiling: bool = False, do_not_save_samples: bool = False, do_not_save_grid: bool = False, extra_generation_params: Dict[Any, Any] = None, overlay_images: Any = None, negative_prompt: str = None, eta: float = None, do_not_reload_embeddings: bool = False, denoising_strength: float = 0, ddim_discretize: str = None, s_churn: float = 0.0, s_tmax: float = None, s_tmin: float = 0.0, s_noise: float = 1.0, override_settings: Dict[str, Any] = None):
+    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt: str = "", styles: List[str] = None, seed: int = -1, subseed: int = -1, subseed_strength: float = 0, seed_resize_from_h: int = -1, seed_resize_from_w: int = -1, seed_enable_extras: bool = True, sampler_index: int = 0, batch_size: int = 1, n_iter: int = 1, steps: int = 50, cfg_scale: float = 7.0, width: int = 512, height: int = 512, restore_faces: bool = False, tiling: bool = False, do_not_save_samples: bool = False, do_not_save_grid: bool = False, extra_generation_params: Dict[Any, Any] = None, overlay_images: Any = None, negative_prompt: str = None, eta: float = None, do_not_reload_embeddings: bool = False, denoising_strength: float = 0, ddim_discretize: str = None, s_churn: float = 0.0, s_tmax: float = None, s_tmin: float = 0.0, s_noise: float = 1.0, script_args: str = None, override_settings: Dict[str, Any] = None):
         self.sd_model = sd_model
-        self.outpath_samples: str = outpath_samples
-        self.outpath_grids: str = outpath_grids
+        self.outpath_samples: str = outpath_samples or opts.outdir_samples or opts.outdir_txt2img_samples
+        self.outpath_grids: str = outpath_grids or opts.outdir_grids or opts.outdir_txt2img_grids
         self.prompt: str = prompt
         self.prompt_for_display: str = None
         self.negative_prompt: str = (negative_prompt or "")
@@ -116,15 +116,14 @@ def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prom
         self.s_tmax = s_tmax or float('inf')  # not representable as a standard ui option
         self.s_noise = s_noise or opts.s_noise
         self.override_settings = {k: v for k, v in (override_settings or {}).items() if k not in shared.restricted_opts}
-
+        self.script_args = json.loads(script_args) if script_args != None else None
+        
         if not seed_enable_extras:
             self.subseed = -1
             self.subseed_strength = 0
             self.seed_resize_from_h = 0
             self.seed_resize_from_w = 0
 
-        self.scripts = None
-        self.script_args = None
         self.all_prompts = None
         self.all_seeds = None
         self.all_subseeds = None

modules/scripts.py

@@ -201,6 +201,22 @@ def __init__(self):
         self.titles = []
         self.infotext_fields = []
 
+    def setup_scripts(self, is_img2img):
+        for script_class, path, basedir in scripts_data:
+            script = script_class()
+            script.filename = path
+
+            visibility = script.show(is_img2img)
+
+            if visibility == AlwaysVisible:
+                self.scripts.append(script)
+                self.alwayson_scripts.append(script)
+                script.alwayson = True
+
+            elif visibility:
+                self.scripts.append(script)
+                self.selectable_scripts.append(script)
+
     def setup_ui(self, is_img2img):
         for script_class, path, basedir in scripts_data:
             script = script_class()

modules/shared.py

@@ -90,6 +90,7 @@
 parser.add_argument("--tls-keyfile", type=str, help="Partially enables TLS, requires --tls-certfile to fully function", default=None)
 parser.add_argument("--tls-certfile", type=str, help="Partially enables TLS, requires --tls-keyfile to fully function", default=None)
 parser.add_argument("--server-name", type=str, help="Sets hostname of server", default=None)
+parser.add_argument("--pureui", action='store_true', help="Pure UI without local inference and progress bar", default=False)
 
 cmd_opts = parser.parse_args()
 restricted_opts = {
@@ -490,6 +491,8 @@ def reorder(self):
 opts = Options()
 if os.path.exists(config_filename):
     opts.load(config_filename)
+if cmd_opts.pureui:
+    opts.show_progressbar = False
 
 sd_upscalers = []
 

webui.py

@@ -8,6 +8,13 @@
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.gzip import GZipMiddleware
 
+import requests
+import json
+import time
+from PIL import Image
+import base64
+import io
+
 from modules.paths import script_path
 
 from modules import devices, sd_samplers, upscaler, extensions, localization
@@ -29,16 +36,93 @@
 
 import modules.ui
 from modules import modelloader
-from modules.shared import cmd_opts
+from modules.shared import cmd_opts,  opts
 import modules.hypernetworks.hypernetwork
 
 queue_lock = threading.Lock()
 server_name = "0.0.0.0" if cmd_opts.listen else cmd_opts.server_name
 
+api_endpoint = os.environ['api_endpoint']
+endpoint_name = os.environ['endpoint_name']
+
 def wrap_queued_call(func):
+    def sagemaker_inference(task, *args, **kwargs):
+        script_args = []
+        for i in range(23, len(args)):
+            script_args.append(args[i])
+ 
+        payload = {
+              "prompt": args[0],
+              "negative_prompt": args[1],
+              "styles": [args[2], args[3]],
+              "steps": args[4],
+              "sampler_index": sd_samplers.samplers[args[5]].name,
+              "restore_faces": args[6],
+              "tiling": args[7],
+              "batch_count": args[8],
+              "batch_size": args[9],
+              "cfg_scale": args[10],
+              "seed": args[11],
+              "subseed": args[12],
+              "subseed_strength": args[13],
+              "seed_resize_from_h": args[14],
+              "seed_resize_from_w": args[15],
+              "seed_checkbox": args[16],
+              "width": args[17],
+              "height": args[18],
+              "enable_hr": args[19],
+              "denoising_strength": args[20],
+              "firstphase_width": args[21],
+              "firstphase_height": args[22],
+              "script_args": json.dumps(script_args),
+              "eta": opts.eta_ddim if sd_samplers.samplers[args[5]].name == 'DDIM' or sd_samplers.samplers[args[5]].name == 'PLMS' else opts.eta_ancestral,
+              "s_churn": opts.s_churn,
+              "s_tmax": None,
+              "s_tmin": opts.s_tmin,
+              "s_noise": opts.s_noise,
+            }
+        inputs = {
+            'task': task,
+            'payload': payload
+        }
+        params = {
+            'endpoint_name': endpoint_name, 
+            'infer_type': 'async'
+        }
+
+        response = requests.post(url=f'{api_endpoint}/inference', params =  params, json = inputs)
+        s3uri = response.text
+        params = {'s3uri': s3uri}
+        start = time.time()
+        while True:
+            response = requests.get(url=f'{api_endpoint}/s3', params = params)
+            text = json.loads(response.text)
+
+            if text['count'] > 0:
+                break
+            else:
+                time.sleep(1)
+        
+        httpuri = text['payload'][0]['httpuri']
+        response = requests.get(url=httpuri)
+        processed = json.loads(response.text)
+        images = []
+        for image in processed['images']:
+            images.append(Image.open(io.BytesIO(base64.b64decode(image.split(',')[1]))))
+        parameters = processed['parameters']
+        info = processed['info']
+        print(f"Time taken: {time.time() - start}s")
+        
+        return images, json.dumps(payload), modules.ui.plaintext_to_html(info)
+
     def f(*args, **kwargs):
-        with queue_lock:
-            res = func(*args, **kwargs)
+        if cmd_opts.pureui and func == modules.txt2img.txt2img:
+            res = sagemaker_inference('text-to-image', *args, **kwargs)
+        elif(cmd_opts.pureui and func == modules.img2img.img2img):
+            res = sagemaker_inference('image-to-image', *args, **kwargs)
+        else:
+            with queue_lock:
+                res = func(*args, **kwargs)
 
         return res
 

webui.sh

@@ -60,19 +60,6 @@ printf "\e[1m\e[32mInstall script for stable-diffusion + Web UI\n"
 printf "\e[1m\e[34mTested on Debian 11 (Bullseye)\e[0m"
 printf "\n%s\n" "${delimiter}"
 
-# Do not run as root
-if [[ $(id -u) -eq 0 ]]
-then
-    printf "\n%s\n" "${delimiter}"
-    printf "\e[1m\e[31mERROR: This script must not be launched as root, aborting...\e[0m"
-    printf "\n%s\n" "${delimiter}"
-    exit 1
-else
-    printf "\n%s\n" "${delimiter}"
-    printf "Running on \e[1m\e[32m%s\e[0m user" "$(whoami)"
-    printf "\n%s\n" "${delimiter}"
-fi
-
 if [[ -d .git ]]
 then
     printf "\n%s\n" "${delimiter}"
@@ -106,6 +93,7 @@ cd "${install_dir}"/ || { printf "\e[1m\e[31mERROR: Can't cd to %s/, aborting...
 if [[ -d "${clone_dir}" ]]
 then
     cd "${clone_dir}"/ || { printf "\e[1m\e[31mERROR: Can't cd to %s/%s/, aborting...\e[0m" "${install_dir}" "${clone_dir}"; exit 1; }
+    "${GIT}" pull
 else
     printf "\n%s\n" "${delimiter}"
     printf "Clone stable-diffusion-webui"