From 57a4a6c5c534ed6b1005d55654d47986af46bbee Mon Sep 17 00:00:00 2001
From: Motto <yjxxtgs@126.COM>
Date: Mon, 13 Mar 2023 19:14:13 +0800
Subject: [PATCH] send to specified controlnet

---
 javascript/main.js | 36 ++++++++++++++++++++++--------------
 scripts/main.py    |  7 ++++++-
 2 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/javascript/main.js b/javascript/main.js
index 529411e..e4d5678 100644
--- a/javascript/main.js
+++ b/javascript/main.js
@@ -117,7 +117,7 @@ function depth_removeBackground() {
     depth_lib_canvas.setBackgroundImage(0, depth_lib_canvas.renderAll.bind(depth_lib_canvas));
 }
 
-function depth_sendImage(){
+function depth_sendImage(index){
     if (depth_lib_canvas.backgroundImage) depth_lib_canvas.backgroundImage.opacity = 0
     depth_lib_canvas.discardActiveObject();
     depth_lib_canvas.renderAll()
@@ -126,19 +126,27 @@ function depth_sendImage(){
         const dt = new DataTransfer();
         dt.items.add(file);
         const list = dt.files
-        depth_gradioApp().querySelector("#txt2img_script_container").querySelectorAll("span.transition").forEach((elem) => {
-            if (elem.previousElementSibling.textContent === "ControlNet"){
-                switch_to_txt2img()
-                elem.className.includes("rotate-90") && elem.parentElement.click();
-                const input = elem.parentElement.parentElement.querySelector("input[type='file']");
-                const button = elem.parentElement.parentElement.querySelector("button[aria-label='Clear']")
-                button && button.click();
-                input.value = "";
-                input.files = list;
-                const event = new Event('change', { 'bubbles': true, "composed": true });
-                input.dispatchEvent(event);
-            }
-        })
+        switch_to_txt2img()
+        const selector = "#txt2img_script_container"
+        const accordion = gradioApp().querySelector(selector).querySelector("#controlnet .transition");
+        if (accordion.classList.contains("rotate-90")) {
+            accordion.click()
+        }
+        const tabs = gradioApp().querySelector(selector).querySelectorAll("#controlnet > div:nth-child(2) > .tabs > .tabitem, #controlnet > div:nth-child(2) > div:not(.tabs)")
+        const tab = tabs[index]
+        if (tab.classList.contains("tabitem")) {
+            tab.parentElement.firstElementChild.querySelector(`:nth-child(${Number(index) + 1})`).click()
+        }
+        const input = tab.querySelector("input[type='file']")
+        try {
+            input.previousElementSibling.previousElementSibling.querySelector("button[aria-label='Clear']").click()
+        } catch (e) {
+            console.error(e)
+        }
+        input.value = "";
+        input.files = list;
+        const event = new Event('change', { 'bubbles': true, "composed": true });
+        input.dispatchEvent(event);
     });
     if (depth_lib_canvas.backgroundImage) depth_lib_canvas.backgroundImage.opacity = 0.5
     depth_lib_canvas.renderAll()
diff --git a/scripts/main.py b/scripts/main.py
index 9a53519..4d4fc6d 100644
--- a/scripts/main.py
+++ b/scripts/main.py
@@ -6,6 +6,7 @@
 
 import modules.scripts as scripts
 from modules import script_callbacks
+from modules.shared import opts
 
 from basicsr.utils.download_util import load_file_from_url
 
@@ -55,6 +56,10 @@ def on_ui_tabs():
         with gr.Row():
           png_output = gr.Button(value="Save PNG")
           send_output = gr.Button(value="Send to ControlNet")
+          control_net_max_models_num = 0
+          if hasattr(opts, 'control_net_max_models_num'):
+            control_net_max_models_num = opts.control_net_max_models_num
+          select_target_index = gr.Dropdown([str(i) for i in range(control_net_max_models_num)], label="Send to", value="0", interactive=True, visible=(control_net_max_models_num > 1))
 
 
     width.change(None, [width, height], None, _js="(w, h) => {depth_resizeCanvas(w, h)}")
@@ -66,7 +71,7 @@ def on_ui_tabs():
     bg_remove.click(None, [], None, _js="depth_removeBackground")
     add.click(None, [png_input_area], None, _js="(path) => {depth_addImg(path)}")
     remove.click(None, [], None, _js="depth_removeSelection")
-    send_output.click(None, [], None, _js="depth_sendImage")
+    send_output.click(None, select_target_index, None, _js="(i) => {depth_sendImage(i)}")
     reset_btn.click(None, [], None, _js="depth_resetCanvas")
 
   return [(depth_lib, "Depth Library", "depth_lib")]