From e89e0682287243c7c14b3d4fe1a8e9e82f68d40d Mon Sep 17 00:00:00 2001
From: Terry Jia <terryjia88@gmail.com>
Date: Sun, 6 Aug 2023 09:56:15 -0400
Subject: [PATCH] move "send to" button to current place

---
 javascript/canvas-editor-import.js   | 18 +++++++++++++++++-
 javascript/lazyload/canvas-editor.js |  8 ++------
 scripts/main.py                      |  6 +-----
 3 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/javascript/canvas-editor-import.js b/javascript/canvas-editor-import.js
index d112559..e202575 100644
--- a/javascript/canvas-editor-import.js
+++ b/javascript/canvas-editor-import.js
@@ -19,7 +19,7 @@
         globalThis.canvasEditor.import = async () => {
             const polotno = await import(`${base_path}/polotno.bundle.js`);
 
-            return { polotno };
+            return {polotno};
         };
 
         if (!globalThis.canvasEditor.imports) {
@@ -38,5 +38,21 @@
     onUiLoaded(function () {
         canvasEditorImport = gradioApp().querySelector('#canvas-editor-import');
         load(canvasEditorImport);
+
+        createSendToCanvasEditorButton("image_buttons_txt2img", window.txt2img_gallery);
+        createSendToCanvasEditorButton("image_buttons_img2img", window.img2img_gallery);
+        createSendToCanvasEditorButton("image_buttons_extras", window.extras_gallery);
+
+        function createSendToCanvasEditorButton(queryId, gallery) {
+            const existingButton = gradioApp().querySelector(`#${queryId} button`);
+            const newButton = existingButton.cloneNode(true);
+            newButton.style.display = "flex";
+            newButton.id = `${queryId}_send_to_canvasEditor`;
+            newButton.textContent = "Send to CanvasEditor";
+            newButton.addEventListener("click", () => sendImageToCanvasEditor(gallery));
+            gradioApp().querySelector(`#${queryId}`).appendChild(newButton);
+        }
     });
+
+
 })();
\ No newline at end of file
diff --git a/javascript/lazyload/canvas-editor.js b/javascript/lazyload/canvas-editor.js
index cfe5977..5892874 100644
--- a/javascript/lazyload/canvas-editor.js
+++ b/javascript/lazyload/canvas-editor.js
@@ -112,12 +112,8 @@ let _r = 0;
         return !(displayValue === 'none');
     }
 
-    window.sendImageToCanvasEditor = function () {
-        const gallerySelector = isTxt2Img() ? '#txt2img_gallery' : '#img2img_gallery';
-
-        const txt2imgGallery = gradioApp().querySelector(gallerySelector);
-
-        const img = txt2imgGallery.querySelector(".preview img");
+    window.sendImageToCanvasEditor = function (gallery) {
+        const img = gallery.querySelector(".preview img");
 
         if (img) {
             const tabIndex = getCanvasEditorTabIndex();
diff --git a/scripts/main.py b/scripts/main.py
index 093c28f..1b89154 100644
--- a/scripts/main.py
+++ b/scripts/main.py
@@ -74,11 +74,7 @@ class Script(scripts.Script):
         return scripts.AlwaysVisible
 
     def ui(self, is_img2img):
-        send_canvas_editor = gr.Button(value="Send to Canvas Editor")
-
-        send_canvas_editor.click(None, send_canvas_editor, None, _js="sendImageToCanvasEditor")
-
-        return [send_canvas_editor]
+        return []
 
 
 def wrap_api(fn):