From cda2f0a1620c3b49bb3408c30796160ed29bc87d Mon Sep 17 00:00:00 2001
From: AUTOMATIC1111 <16777216c@gmail.com>
Date: Mon, 14 Aug 2023 08:49:39 +0300
Subject: [PATCH] make on_before_component/on_after_component possible earlier

---
 modules/scripts.py | 35 +++++++++++++++++++++++++----------
 1 file changed, 25 insertions(+), 10 deletions(-)

diff --git a/modules/scripts.py b/modules/scripts.py
index d4a9da94..cbdac2b5 100644
--- a/modules/scripts.py
+++ b/modules/scripts.py
@@ -239,6 +239,8 @@ class Script:
         """
         Calls callback before a component is created. The callback function is called with a single argument of type OnComponent.
 
+        May be called in show() or ui() - but it may be too late in latter as some components may already be created.
+
         This function is an alternative to before_component in that it also cllows to run before a component is created, but
         it doesn't require to be called for every created component - just for the one you need.
         """
@@ -445,6 +447,28 @@ class ScriptRunner:
                 self.scripts.append(script)
                 self.selectable_scripts.append(script)
 
+        self.apply_on_before_component_callbacks()
+
+    def apply_on_before_component_callbacks(self):
+        for script in self.scripts:
+            on_before = script.on_before_component_elem_id or []
+            on_after = script.on_after_component_elem_id or []
+
+            for elem_id, callback in on_before:
+                if elem_id not in self.on_before_component_elem_id:
+                    self.on_before_component_elem_id[elem_id] = []
+
+                self.on_before_component_elem_id[elem_id].append((callback, script))
+
+            for elem_id, callback in on_after:
+                if elem_id not in self.on_after_component_elem_id:
+                    self.on_after_component_elem_id[elem_id] = []
+
+                self.on_after_component_elem_id[elem_id].append((callback, script))
+
+            on_before.clear()
+            on_after.clear()
+
     def create_script_ui(self, script):
         import modules.api.models as api_models
 
@@ -555,16 +579,7 @@ class ScriptRunner:
         self.infotext_fields.append((dropdown, lambda x: gr.update(value=x.get('Script', 'None'))))
         self.infotext_fields.extend([(script.group, onload_script_visibility) for script in self.selectable_scripts])
 
-        for script in self.scripts:
-            for elem_id, callback in script.on_before_component_elem_id or []:
-                items = self.on_before_component_elem_id.get(elem_id, [])
-                items.append((callback, script))
-                self.on_before_component_elem_id[elem_id] = items
-
-            for elem_id, callback in script.on_after_component_elem_id or []:
-                items = self.on_after_component_elem_id.get(elem_id, [])
-                items.append((callback, script))
-                self.on_after_component_elem_id[elem_id] = items
+        self.apply_on_before_component_callbacks()
 
         return self.inputs