215 lines
7 KiB
Python
215 lines
7 KiB
Python
import os
|
|
import sys
|
|
import traceback
|
|
|
|
import modules.ui as ui
|
|
import gradio as gr
|
|
|
|
from modules.processing import StableDiffusionProcessing
|
|
from modules import shared
|
|
|
|
class Script:
|
|
filename = None
|
|
args_from = None
|
|
args_to = None
|
|
|
|
# The title of the script. This is what will be displayed in the dropdown menu.
|
|
def title(self):
|
|
raise NotImplementedError()
|
|
|
|
# How the script is displayed in the UI. See https://gradio.app/docs/#components
|
|
# for the different UI components you can use and how to create them.
|
|
# Most UI components can return a value, such as a boolean for a checkbox.
|
|
# The returned values are passed to the run method as parameters.
|
|
def ui(self, is_img2img):
|
|
pass
|
|
|
|
# Determines when the script should be shown in the dropdown menu via the
|
|
# returned value. As an example:
|
|
# is_img2img is True if the current tab is img2img, and False if it is txt2img.
|
|
# Thus, return is_img2img to only show the script on the img2img tab.
|
|
def show(self, is_img2img):
|
|
return True
|
|
|
|
# This is where the additional processing is implemented. The parameters include
|
|
# self, the model object "p" (a StableDiffusionProcessing class, see
|
|
# processing.py), and the parameters returned by the ui method.
|
|
# Custom functions can be defined here, and additional libraries can be imported
|
|
# to be used in processing. The return value should be a Processed object, which is
|
|
# what is returned by the process_images method.
|
|
def run(self, *args):
|
|
raise NotImplementedError()
|
|
|
|
# The description method is currently unused.
|
|
# To add a description that appears when hovering over the title, amend the "titles"
|
|
# dict in script.js to include the script title (returned by title) as a key, and
|
|
# your description as the value.
|
|
def describe(self):
|
|
return ""
|
|
|
|
|
|
scripts_data = []
|
|
|
|
|
|
def load_scripts(basedir):
|
|
if not os.path.exists(basedir):
|
|
return
|
|
|
|
for filename in sorted(os.listdir(basedir)):
|
|
path = os.path.join(basedir, filename)
|
|
|
|
if os.path.splitext(path)[1].lower() != '.py':
|
|
continue
|
|
|
|
if not os.path.isfile(path):
|
|
continue
|
|
|
|
try:
|
|
with open(path, "r", encoding="utf8") as file:
|
|
text = file.read()
|
|
|
|
from types import ModuleType
|
|
compiled = compile(text, path, 'exec')
|
|
module = ModuleType(filename)
|
|
exec(compiled, module.__dict__)
|
|
|
|
for key, script_class in module.__dict__.items():
|
|
if type(script_class) == type and issubclass(script_class, Script):
|
|
scripts_data.append((script_class, path))
|
|
|
|
except Exception:
|
|
print(f"Error loading script: {filename}", file=sys.stderr)
|
|
print(traceback.format_exc(), file=sys.stderr)
|
|
|
|
|
|
def wrap_call(func, filename, funcname, *args, default=None, **kwargs):
|
|
try:
|
|
res = func(*args, **kwargs)
|
|
return res
|
|
except Exception:
|
|
print(f"Error calling: {filename}/{funcname}", file=sys.stderr)
|
|
print(traceback.format_exc(), file=sys.stderr)
|
|
|
|
return default
|
|
|
|
|
|
class ScriptRunner:
|
|
def __init__(self):
|
|
self.scripts = []
|
|
self.titles = []
|
|
|
|
def setup_ui(self, is_img2img):
|
|
for script_class, path in scripts_data:
|
|
script = script_class()
|
|
script.filename = path
|
|
|
|
if not script.show(is_img2img):
|
|
continue
|
|
|
|
self.scripts.append(script)
|
|
|
|
self.titles = [wrap_call(script.title, script.filename, "title") or f"{script.filename} [error]" for script in self.scripts]
|
|
|
|
dropdown = gr.Dropdown(label="Script", choices=["None"] + self.titles, value="None", type="index")
|
|
dropdown.save_to_config = True
|
|
inputs = [dropdown]
|
|
|
|
for script in self.scripts:
|
|
script.args_from = len(inputs)
|
|
script.args_to = len(inputs)
|
|
|
|
controls = wrap_call(script.ui, script.filename, "ui", is_img2img)
|
|
|
|
if controls is None:
|
|
continue
|
|
|
|
for control in controls:
|
|
control.custom_script_source = os.path.basename(script.filename)
|
|
control.visible = False
|
|
|
|
inputs += controls
|
|
script.args_to = len(inputs)
|
|
|
|
def select_script(script_index):
|
|
if 0 < script_index <= len(self.scripts):
|
|
script = self.scripts[script_index-1]
|
|
args_from = script.args_from
|
|
args_to = script.args_to
|
|
else:
|
|
args_from = 0
|
|
args_to = 0
|
|
|
|
return [ui.gr_show(True if i == 0 else args_from <= i < args_to) for i in range(len(inputs))]
|
|
|
|
def init_field(title):
|
|
if title == 'None':
|
|
return
|
|
script_index = self.titles.index(title)
|
|
script = self.scripts[script_index]
|
|
for i in range(script.args_from, script.args_to):
|
|
inputs[i].visible = True
|
|
|
|
dropdown.init_field = init_field
|
|
dropdown.change(
|
|
fn=select_script,
|
|
inputs=[dropdown],
|
|
outputs=inputs
|
|
)
|
|
|
|
return inputs
|
|
|
|
def run(self, p: StableDiffusionProcessing, *args):
|
|
script_index = args[0]
|
|
|
|
if script_index == 0:
|
|
return None
|
|
|
|
script = self.scripts[script_index-1]
|
|
|
|
if script is None:
|
|
return None
|
|
|
|
script_args = args[script.args_from:script.args_to]
|
|
processed = script.run(p, *script_args)
|
|
|
|
shared.total_tqdm.clear()
|
|
|
|
return processed
|
|
|
|
def reload_sources(self):
|
|
for si, script in list(enumerate(self.scripts)):
|
|
with open(script.filename, "r", encoding="utf8") as file:
|
|
args_from = script.args_from
|
|
args_to = script.args_to
|
|
filename = script.filename
|
|
text = file.read()
|
|
|
|
from types import ModuleType
|
|
|
|
compiled = compile(text, filename, 'exec')
|
|
module = ModuleType(script.filename)
|
|
exec(compiled, module.__dict__)
|
|
|
|
for key, script_class in module.__dict__.items():
|
|
if type(script_class) == type and issubclass(script_class, Script):
|
|
self.scripts[si] = script_class()
|
|
self.scripts[si].filename = filename
|
|
self.scripts[si].args_from = args_from
|
|
self.scripts[si].args_to = args_to
|
|
|
|
scripts_txt2img = ScriptRunner()
|
|
scripts_img2img = ScriptRunner()
|
|
|
|
def reload_script_body_only():
|
|
scripts_txt2img.reload_sources()
|
|
scripts_img2img.reload_sources()
|
|
|
|
|
|
def reload_scripts(basedir):
|
|
global scripts_txt2img, scripts_img2img
|
|
|
|
scripts_data.clear()
|
|
load_scripts(basedir)
|
|
|
|
scripts_txt2img = ScriptRunner()
|
|
scripts_img2img = ScriptRunner()
|