custom unet support

This commit is contained in:
AUTOMATIC
2023-05-27 15:47:33 +03:00
parent a6e653be26
commit 339b531570
8 changed files with 148 additions and 8 deletions

View File

@@ -111,6 +111,7 @@ callback_map = dict(
callbacks_before_ui=[],
callbacks_on_reload=[],
callbacks_list_optimizers=[],
callbacks_list_unets=[],
)
@@ -271,6 +272,18 @@ def list_optimizers_callback():
return res
def list_unets_callback():
res = []
for c in callback_map['callbacks_list_unets']:
try:
c.callback(res)
except Exception:
report_exception(c, 'list_unets')
return res
def add_callback(callbacks, fun):
stack = [x for x in inspect.stack() if x.filename != __file__]
filename = stack[0].filename if len(stack) > 0 else 'unknown file'
@@ -430,3 +443,10 @@ def on_list_optimizers(callback):
to it."""
add_callback(callback_map['callbacks_list_optimizers'], callback)
def on_list_unets(callback):
"""register a function to be called when UI is making a list of alternative options for unet.
The function will be called with one argument, a list, and shall add objects of type modules.sd_unet.SdUnetOption to it."""
add_callback(callback_map['callbacks_list_unets'], callback)