Commit 79d57d02 authored by space-nuko's avatar space-nuko

Improve custom code extension

- Uses `gr.Code` component
- Includes example
- Can return out of body
parent 22bcc7be
import modules.scripts as scripts import modules.scripts as scripts
import gradio as gr import gradio as gr
import ast
import copy
from modules.processing import Processed from modules.processing import Processed
from modules.shared import opts, cmd_opts, state from modules.shared import opts, cmd_opts, state
def convertExpr2Expression(expr):
expr.lineno = 0
expr.col_offset = 0
result = ast.Expression(expr.value, lineno=0, col_offset = 0)
return result
def exec_with_return(code, module):
"""
like exec() but can return values
https://stackoverflow.com/a/52361938/5862977
"""
code_ast = ast.parse(code)
init_ast = copy.deepcopy(code_ast)
init_ast.body = code_ast.body[:-1]
last_ast = copy.deepcopy(code_ast)
last_ast.body = code_ast.body[-1:]
exec(compile(init_ast, "<ast>", "exec"), module.__dict__)
if type(last_ast.body[0]) == ast.Expr:
return eval(compile(convertExpr2Expression(last_ast.body[0]), "<ast>", "eval"), module.__dict__)
else:
exec(compile(last_ast, "<ast>", "exec"), module.__dict__)
class Script(scripts.Script): class Script(scripts.Script):
def title(self): def title(self):
...@@ -13,12 +44,23 @@ class Script(scripts.Script): ...@@ -13,12 +44,23 @@ class Script(scripts.Script):
return cmd_opts.allow_code return cmd_opts.allow_code
def ui(self, is_img2img): def ui(self, is_img2img):
code = gr.Textbox(label="Python code", lines=1, elem_id=self.elem_id("code")) example = """from modules.processing import process_images
p.width = 768
p.height = 768
p.batch_size = 2
p.steps = 10
return process_images(p)
"""
return [code] code = gr.Code(value=example, language="python", label="Python code", elem_id=self.elem_id("code"))
indent_level = gr.Number(label='Indent level', value=2, precision=0, elem_id=self.elem_id("indent_level"))
return [code, indent_level]
def run(self, p, code): def run(self, p, code, indent_level):
assert cmd_opts.allow_code, '--allow-code option must be enabled' assert cmd_opts.allow_code, '--allow-code option must be enabled'
display_result_data = [[], -1, ""] display_result_data = [[], -1, ""]
...@@ -29,13 +71,20 @@ class Script(scripts.Script): ...@@ -29,13 +71,20 @@ class Script(scripts.Script):
display_result_data[2] = i display_result_data[2] = i
from types import ModuleType from types import ModuleType
compiled = compile(code, '', 'exec')
module = ModuleType("testmodule") module = ModuleType("testmodule")
module.__dict__.update(globals()) module.__dict__.update(globals())
module.p = p module.p = p
module.display = display module.display = display
exec(compiled, module.__dict__)
return Processed(p, *display_result_data) indent = " " * indent_level
indented = code.replace('\n', '\n' + indent)
body = f"""def __webuitemp__():
{indent}{indented}
__webuitemp__()"""
result = exec_with_return(body, module)
if isinstance(result, Processed):
return result
return Processed(p, *display_result_data)
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment