1 files changed, 40 insertions, 0 deletions
diff --git a/scripts/custom_code.py b/scripts/custom_code.py
new file mode 100644
index 00000000..b359050a
--- /dev/null
+++ b/scripts/custom_code.py
@@ -0,0 +1,40 @@
+import modules.scripts as scripts
+import gradio as gr
+
+from modules.processing import Processed
+from modules.shared import opts, cmd_opts, state
+
+
+class Script(scripts.Script):
+    def title(self):
+        return "Custom code"
+
+    def enabled(self):
+        return cmd_opts.allow_code
+
+    def ui(self, is_img2img):
+        code = gr.Textbox(label="Python code", visible=False, lines=1)
+
+        return [code]
+
+    def run(self, p, code):
+        if not cmd_opts.allow_code:
+            return
+
+        display_result_data = [[], -1, ""]
+
+        def display(imgs, s=display_result_data[1], i=display_result_data[2]):
+            display_result_data[0] = imgs
+            display_result_data[1] = s
+            display_result_data[2] = i
+
+        from types import ModuleType
+        compiled = compile(code, '', 'exec')
+        module = ModuleType("testmodule")
+        module.__dict__.update(globals())
+        module.p = p
+        module.display = display
+        exec(compiled, module.__dict__)
+
+        return Processed(p, *display_result_data)
+