From 2bf7116f011868ab8675aa6fa94e215f1b813d77 Mon Sep 17 00:00:00 2001
From: ManInDark <61268856+ManInDark@users.noreply.github.com>
Date: Wed, 9 Aug 2023 22:50:03 +0200
Subject: [PATCH] Added Gallery search for prompt & model

---
 ui/easydiffusion/bucket_manager.py | 11 +++++++----
 ui/index.html                      |  6 +++++-
 ui/media/css/main.css              | 13 ++++++++++++-
 ui/media/js/main.js                | 13 ++++++++++++-
 4 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/ui/easydiffusion/bucket_manager.py b/ui/easydiffusion/bucket_manager.py
index 0d72ed06..43b5f147 100644
--- a/ui/easydiffusion/bucket_manager.py
+++ b/ui/easydiffusion/bucket_manager.py
@@ -100,11 +100,14 @@ def init():
             raise HTTPException(status_code=404, detail="Image not found")
     
     @server_api.get("/all_images")
-    def get_all_images(db: Session = Depends(get_db)):
+    def get_all_images(prompt: str = "", model: str = "", db: Session = Depends(get_db)):
         from easydiffusion.easydb.mappings import GalleryImage
-        images = db.query(GalleryImage).all()
-        return images
-
+        images = db.query(GalleryImage)
+        if prompt != "":
+            images = images.filter(GalleryImage.path.like("%"+prompt+"%"))
+        if model != "":
+            images = images.filter(GalleryImage.use_stable_diffusion_model.like("%"+model+"%"))
+        return images.all()
 
 def get_filename_from_url(url):
     path = urlparse(url).path
diff --git a/ui/index.html b/ui/index.html
index d5803fcc..9aeb7ab2 100644
--- a/ui/index.html
+++ b/ui/index.html
@@ -517,7 +517,11 @@
             </div>
         </div>
         <div id="tab-content-gallery" class="tab-content">
-            <button class="primaryButton" onclick="refreshGallery()">Refresh</button>
+            <div id="gallery-search">
+                <textarea id="gallery-prompt-search" onkeydown="gallery_keyDown_handler(event)" placeholder="Search for a prompt..."></textarea>
+                <textarea id="gallery-model-search" onkeydown="gallery_keyDown_handler(event)" placeholder="Search for a model..."></textarea>
+                <button class="primaryButton" onclick="refreshGallery()">Refresh</button>
+            </div>
             <div class="gallery">
                 <div class="gallery-container" id="imagecontainer"></div>
             </div>
diff --git a/ui/media/css/main.css b/ui/media/css/main.css
index e0d1d02d..d41d6f17 100644
--- a/ui/media/css/main.css
+++ b/ui/media/css/main.css
@@ -1897,6 +1897,17 @@ div#enlarge-buttons {
   font-family: sans-serif;
 }
 
+#gallery-search {
+    display: flex;
+    justify-content: center;
+    align-items: center;
+    flex-direction: row;
+}
+
+#gallery-search>* {
+    margin: 0.5em;
+}
+
 .gallery-container {
   columns: 5 ;
   column-gap: 1.5rem;
@@ -1924,6 +1935,6 @@ div#enlarge-buttons {
 }
 
 
-#tab-content-gallery>button {
+#tab-content-gallery>* {
     margin: 8px;
 }
diff --git a/ui/media/js/main.js b/ui/media/js/main.js
index a4476252..5beb094d 100644
--- a/ui/media/js/main.js
+++ b/ui/media/js/main.js
@@ -3097,8 +3097,12 @@ function galleryImage(item) {
 
 function refreshGallery() {
     let container = document.getElementById("imagecontainer")
+    let promptsearchfield = document.getElementById("gallery-prompt-search").value
+    let promptsearch = promptsearchfield.length > 0 ? "prompt=" + promptsearchfield  + "&" : ""
+    let modelsearchfield = document.getElementById("gallery-model-search").value
+    let modelsearch = modelsearchfield.length > 0 ? "model=" + modelsearchfield  + "&" : ""
     container.innerHTML=""
-    fetch('/all_images')
+    fetch('/all_images?' + promptsearch + modelsearch)
         .then(response => response.json())
         .then(json => {
             console.log(json)
@@ -3107,3 +3111,10 @@ function refreshGallery() {
             })
          })
 }
+
+function gallery_keyDown_handler(event) {
+    if (event.key === 'Enter') {
+        event.preventDefault()
+        refreshGallery()
+    }
+}
\ No newline at end of file