Merge branch 'dev' into img2img-batch-png-info

1 files changed, 28 insertions, 7 deletions

diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py
index 40f388a5..77ee55ee 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -1,16 +1,16 @@
-import sys

-import traceback

-from collections import namedtuple

 import inspect

+import os

+from collections import namedtuple

 from typing import Optional, Dict, Any

 

 from fastapi import FastAPI

 from gradio import Blocks

 

+from modules import errors, timer

+

 

 def report_exception(c, job):

-    print(f"Error executing callback {job} for {c.script}", file=sys.stderr)

-    print(traceback.format_exc(), file=sys.stderr)

+    errors.report(f"Error executing callback {job} for {c.script}", exc_info=True)

 

 

 class ImageSaveParams:

@@ -111,6 +111,7 @@ callback_map = dict(
     callbacks_before_ui=[],

     callbacks_on_reload=[],

     callbacks_list_optimizers=[],

+    callbacks_list_unets=[],

 )

 

 

@@ -123,6 +124,7 @@ def app_started_callback(demo: Optional[Blocks], app: FastAPI):
     for c in callback_map['callbacks_app_started']:

         try:

             c.callback(demo, app)

+            timer.startup_timer.record(os.path.basename(c.script))

         except Exception:

             report_exception(c, 'app_started_callback')

 

@@ -271,16 +273,28 @@ 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'

+    filename = stack[0].filename if stack else 'unknown file'

 

     callbacks.append(ScriptCallback(filename, fun))

 

 

 def remove_current_script_callbacks():

     stack = [x for x in inspect.stack() if x.filename != __file__]

-    filename = stack[0].filename if len(stack) > 0 else 'unknown file'

+    filename = stack[0].filename if stack else 'unknown file'

     if filename == 'unknown file':

         return

     for callback_list in callback_map.values():

@@ -430,3 +444,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)