From d38ede71d5330958f4bbac5f99c1be3c146b506a Mon Sep 17 00:00:00 2001
From: noodleanon <122053346+noodleanon@users.noreply.github.com>
Date: Sat, 7 Jan 2023 14:21:31 +0000
Subject: [PATCH] Added script support in txt2img endpoint

---
 modules/api/api.py    | 22 +++++++++++++++++++---
 modules/api/models.py |  2 +-
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/modules/api/api.py b/modules/api/api.py
index aa62a42e..0e8ea263 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -149,6 +149,14 @@ class Api:
         raise HTTPException(status_code=401, detail="Incorrect username or password", headers={"WWW-Authenticate": "Basic"})
 
     def text2imgapi(self, txt2imgreq: StableDiffusionTxt2ImgProcessingAPI):
+        if txt2imgreq.script_name is not None:
+            if scripts.scripts_txt2img.scripts == []:
+                scripts.scripts_txt2img.initialize_scripts(True)
+                ui.create_ui()
+
+            script_idx = script_name_to_index(txt2imgreq.script_name, scripts.scripts_txt2img.selectable_scripts)
+            script = scripts.scripts_txt2img.selectable_scripts[script_idx]
+
         populate = txt2imgreq.copy(update={ # Override __init__ params
             "sampler_name": validate_sampler_name(txt2imgreq.sampler_name or txt2imgreq.sampler_index),
             "do_not_save_samples": True,
@@ -158,11 +166,20 @@ class Api:
         if populate.sampler_name:
             populate.sampler_index = None  # prevent a warning later on
 
+        args = vars(populate)
+        args.pop('script_name', None)
+
         with self.queue_lock:
-            p = StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **vars(populate))
+            p = StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)
 
             shared.state.begin()
-            processed = process_images(p)
+            if 'script' in locals():
+                p.outpath_grids = opts.outdir_txt2img_grids
+                p.outpath_samples = opts.outdir_txt2img_samples
+                p.script_args = [script_idx + 1] + [None] * (script.args_from - 1) + p.script_args
+                processed = scripts.scripts_txt2img.run(p, *p.script_args)
+            else:
+                processed = process_images(p)
             shared.state.end()
 
 
@@ -213,7 +230,6 @@ class Api:
                 processed = scripts.scripts_img2img.run(p, *p.script_args)
             else:
                 processed = process_images(p)
-
             shared.state.end()
 
         b64images = list(map(encode_pil_to_base64, processed.images))
diff --git a/modules/api/models.py b/modules/api/models.py
index c85eb94d..ce43c858 100644
--- a/modules/api/models.py
+++ b/modules/api/models.py
@@ -100,7 +100,7 @@ class PydanticModelGenerator:
 StableDiffusionTxt2ImgProcessingAPI = PydanticModelGenerator(
     "StableDiffusionProcessingTxt2Img",
     StableDiffusionProcessingTxt2Img,
-    [{"key": "sampler_index", "type": str, "default": "Euler"}]
+    [{"key": "sampler_index", "type": str, "default": "Euler"}, {"key": "script_name", "type": str, "default": None}, {"key": "script_args", "type": list, "default": []}]
 ).generate_model()
 
 StableDiffusionImg2ImgProcessingAPI = PydanticModelGenerator(