stable-diffusion-webui/scripts/custom_code.py

42 lines
1.1 KiB
Python
Raw Normal View History

2022-09-03 10:32:45 -06:00
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):
2022-09-03 10:32:45 -06:00
def title(self):
return "Custom code"
2022-09-04 04:08:06 -06:00
def show(self, is_img2img):
2022-09-03 10:32:45 -06:00
return cmd_opts.allow_code
def ui(self, is_img2img):
code = gr.Textbox(label="Python code", lines=1)
2022-09-03 10:32:45 -06:00
return [code]
2022-09-03 10:32:45 -06:00
def run(self, p, code):
2022-09-04 04:08:06 -06:00
assert cmd_opts.allow_code, '--allow-code option must be enabled'
2022-09-03 10:32:45 -06:00
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)