From 9c2a7f1e8bafcb59e566bf568fdefe1be95905fe Mon Sep 17 00:00:00 2001
From: w-e-w <40751091+w-e-w@users.noreply.github.com>
Date: Mon, 19 Jun 2023 15:37:20 +0900
Subject: [PATCH] add callback after_extra_networks_activate

---
 modules/extra_networks.py |  3 +++
 modules/scripts.py        | 23 +++++++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/modules/extra_networks.py b/modules/extra_networks.py
index 1f093df27..41799b0a9 100644
--- a/modules/extra_networks.py
+++ b/modules/extra_networks.py
@@ -103,6 +103,9 @@ def activate(p, extra_network_data):
         except Exception as e:
             errors.display(e, f"activating extra network {extra_network_name}")
 
+    if p.scripts is not None:
+        p.scripts.after_extra_networks_activate(p, batch_number=p.iteration, prompts=p.prompts, seeds=p.seeds, subseeds=p.subseeds, extra_network_data=extra_network_data)
+
 
 def deactivate(p, extra_network_data):
     """call deactivate for extra networks in extra_network_data in specified order, then call
diff --git a/modules/scripts.py b/modules/scripts.py
index 99bf836af..340f1480b 100644
--- a/modules/scripts.py
+++ b/modules/scripts.py
@@ -116,6 +116,21 @@ class Script:
 
         pass
 
+    def after_extra_networks_activate(self, p, *args, **kwargs):
+        """
+        Calledafter extra networks activation, before conds calculation
+        allow modification of the network after extra networks activation been applied
+        won't be call if p.disable_extra_networks
+
+        **kwargs will have those items:
+          - batch_number - index of current batch, from 0 to number of batches-1
+          - prompts - list of prompts for current batch; you can change contents of this list but changing the number of entries will likely break things
+          - seeds - list of seeds for current batch
+          - subseeds - list of subseeds for current batch
+          - extra_network_data - list of ExtraNetworkParams for current stage
+        """
+        pass
+
     def process_batch(self, p, *args, **kwargs):
         """
         Same as process(), but called for every batch.
@@ -483,6 +498,14 @@ class ScriptRunner:
             except Exception:
                 errors.report(f"Error running before_process_batch: {script.filename}", exc_info=True)
 
+    def after_extra_networks_activate(self, p, **kwargs):
+        for script in self.alwayson_scripts:
+            try:
+                script_args = p.script_args[script.args_from:script.args_to]
+                script.after_extra_networks_activate(p, *script_args, **kwargs)
+            except Exception:
+                errors.report(f"Error running after_extra_networks_activate: {script.filename}", exc_info=True)
+
     def process_batch(self, p, **kwargs):
         for script in self.alwayson_scripts:
             try: