sdgsdggds commited on
Commit
fe3fb2f
·
1 Parent(s): af4dcd7

Upload 7 files

Browse files
Files changed (7) hide show
  1. .gitignore +6 -0
  2. README.md +61 -0
  3. install.py +8 -0
  4. javascript/image_browser.js +770 -0
  5. scripts/image_browser.py +1715 -0
  6. scripts/wib/wib_db.py +888 -0
  7. style.css +15 -0
.gitignore ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ .DS_Store
2
+ path_recorder.txt
3
+ __pycache__
4
+ *.json
5
+ *.sqlite3
6
+ *.log
README.md ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ## stable-diffusion-webui-images-browser
2
+
3
+ A custom extension for [AUTOMATIC1111/stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui).
4
+
5
+ This is an image browser for browsing past generated pictures, view their generated informations, send that information to txt2img, img2img and others, collect images to your "favorites" folder and delete the images you no longer need.
6
+
7
+ ## Installation
8
+
9
+ The extension can be installed directly from within the **Extensions** tab within the Webui.
10
+
11
+ You can also install it manually by running the following command from within the webui directory:
12
+
13
+ git clone https://github.com/AlUlkesh/stable-diffusion-webui-images-browser/ extensions/stable-diffusion-webui-images-browser
14
+
15
+ and restart your stable-diffusion-webui, then you can see the new tab "Image Browser".
16
+
17
+ Please be aware that when scanning a directory for the first time, the png-cache will be built. This can take several minutes, depending on the amount of images.
18
+
19
+ ## Recent updates
20
+ - Image Reward scoring
21
+ - Size tooltip for thumbnails
22
+ - Optimized images in the thumbnail interface
23
+ - Send to ControlNet
24
+ - Hidable UI components
25
+ - Send to openOutpaint
26
+ - Regex search
27
+ - Maximum aesthetic_score filter
28
+ - Save ranking to EXIF option
29
+ - Maintenance tab
30
+ - Custom tabs
31
+ - Copy/Move to directory
32
+ - Keybindings
33
+ - Additional sorting and filtering by EXIF data including .txt file information
34
+ - Recyle bin option
35
+ - Add/Remove from saved directories, via buttons
36
+ - New dropdown with subdirs
37
+ - Option to not show the images from subdirs
38
+ - Refresh button
39
+ - Sort order
40
+ - View and save favorites with individual folder depth
41
+ - Now also supports jpg
42
+
43
+ Please also check the [discussions](https://github.com/AlUlkesh/stable-diffusion-webui-images-browser/discussions) for major update information.
44
+
45
+ ## Keybindings
46
+ | Key | Explanation |
47
+ |---------|-------------|
48
+ | `0-5` | Ranks the current image, with 0 being the last option (None) |
49
+ | `F` | Adds the current image to Favorites |
50
+ | `R` | Refreshes the image gallery |
51
+ | `Delete` | Deletes the current image |
52
+ | `Ctrl + Arrow Left` | Goes to the previous page of images |
53
+ | `Ctrl + Arrow Right` | Goes to the next page of images |
54
+
55
+ (Ctrl can be changed in settings)
56
+
57
+ ## Credit
58
+
59
+ Credit goes to the original maintainer of this extension: https://github.com/yfszzx and to major contributors https://github.com/Klace and https://github.com/EllangoK
60
+
61
+ Image Reward: https://github.com/THUDM/ImageReward
install.py ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ import launch
2
+
3
+ if not launch.is_installed("send2trash"):
4
+ launch.run_pip("install Send2Trash", "Send2Trash requirement for image browser")
5
+
6
+ # temporarily deactivated
7
+ #if not launch.is_installed("ImageReward"):
8
+ #launch.run_pip("install image-reward", "ImageReward requirement for image browser")
javascript/image_browser.js ADDED
@@ -0,0 +1,770 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ let image_browser_state = "free"
2
+ let image_browser_webui_ready = false
3
+ let image_browser_started = false
4
+ let image_browser_console_log = ""
5
+ let image_browser_debug = false
6
+ let image_browser_img_show_in_progress = false
7
+
8
+ function image_browser_delay(ms){return new Promise(resolve => setTimeout(resolve, ms))}
9
+
10
+ onUiLoaded(image_browser_start_it_up)
11
+
12
+ async function image_browser_wait_for_webui() {
13
+ if (image_browser_debug) console.log("image_browser_wait_for_webui:start")
14
+ await image_browser_delay(500)
15
+ sd_model = gradioApp().getElementById("setting_sd_model_checkpoint")
16
+ if (!sd_model.querySelector(".eta-bar")) {
17
+ image_browser_webui_ready = true
18
+ image_browser_start()
19
+ } else {
20
+ // Set timeout for MutationObserver
21
+ const startTime = Date.now()
22
+ // 40 seconds in milliseconds
23
+ const timeout = 40000
24
+ const webuiObserver = new MutationObserver(function(mutationsList) {
25
+ if (image_browser_debug) console.log("webuiObserver:start")
26
+ let found = false
27
+ outerLoop: for (let i = 0; i < mutationsList.length; i++) {
28
+ let mutation = mutationsList[i];
29
+ if (mutation.type === "childList") {
30
+ for (let j = 0; j < mutation.removedNodes.length; j++) {
31
+ let node = mutation.removedNodes[j];
32
+ if (node.nodeType === Node.ELEMENT_NODE && node.classList.contains("eta-bar")) {
33
+ found = true
34
+ break outerLoop;
35
+ }
36
+ }
37
+ }
38
+ }
39
+ if (found || (Date.now() - startTime > timeout)) {
40
+ image_browser_webui_ready = true
41
+ webuiObserver.disconnect()
42
+ if (image_browser_debug) console.log("webuiObserver:end")
43
+ image_browser_start()
44
+ }
45
+ })
46
+ webuiObserver.observe(gradioApp(), { childList:true, subtree:true })
47
+ }
48
+ if (image_browser_debug) console.log("image_browser_wait_for_webui:end")
49
+ }
50
+
51
+ async function image_browser_start_it_up() {
52
+ if (image_browser_debug) console.log("image_browser_start_it_up:start")
53
+ container = gradioApp().getElementById("image_browser_tabs_container")
54
+ let controls = container.querySelectorAll('[id*="_control_"]')
55
+ controls.forEach(function(control) {
56
+ control.style.pointerEvents = "none"
57
+ control.style.cursor = "not-allowed"
58
+ control.style.opacity = "0.65"
59
+ })
60
+ let warnings = container.querySelectorAll('[id*="_warning_box"]')
61
+ warnings.forEach(function(warning) {
62
+ warning.innerHTML = '<p style="font-weight: bold;">Waiting for webui...'
63
+ })
64
+
65
+ image_browser_wait_for_webui()
66
+ if (image_browser_debug) console.log("image_browser_start_it_up:end")
67
+ }
68
+
69
+ async function image_browser_lock(reason) {
70
+ if (image_browser_debug) console.log("image_browser_lock:start")
71
+ // Wait until lock removed
72
+ let i = 0
73
+ while (image_browser_state != "free") {
74
+ await image_browser_delay(200)
75
+ i = i + 1
76
+ if (i === 150) {
77
+ throw new Error("Still locked after 30 seconds. Please Reload UI.")
78
+ }
79
+ }
80
+ // Lock
81
+ image_browser_state = reason
82
+ if (image_browser_debug) console.log("image_browser_lock:end")
83
+ }
84
+
85
+ async function image_browser_unlock() {
86
+ if (image_browser_debug) console.log("image_browser_unlock:start")
87
+ image_browser_state = "free"
88
+ if (image_browser_debug) console.log("image_browser_unlock:end")
89
+ }
90
+
91
+ const image_browser_click_image = async function() {
92
+ if (image_browser_debug) console.log("image_browser_click_image:start")
93
+ await image_browser_lock("image_browser_click_image")
94
+ const tab_base_tag = image_browser_current_tab()
95
+ const container = gradioApp().getElementById(tab_base_tag + "_image_browser_container")
96
+ let child = this
97
+ let index = 0
98
+ while((child = child.previousSibling) != null) {
99
+ index = index + 1
100
+ }
101
+ const set_btn = container.querySelector(".image_browser_set_index")
102
+ let curr_idx
103
+ try {
104
+ curr_idx = set_btn.getAttribute("img_index")
105
+ } catch (e) {
106
+ curr_idx = -1
107
+ }
108
+ if (curr_idx != index) {
109
+ set_btn.setAttribute("img_index", index)
110
+ }
111
+ await image_browser_unlock()
112
+ set_btn.click()
113
+ if (image_browser_debug) console.log("image_browser_click_image:end")
114
+ }
115
+
116
+ async function image_browser_get_current_img(tab_base_tag, img_index, page_index, filenames, turn_page_switch, image_gallery) {
117
+ if (image_browser_debug) console.log("image_browser_get_current_img:start")
118
+ await image_browser_lock("image_browser_get_current_img")
119
+ img_index = gradioApp().getElementById(tab_base_tag + '_image_browser_set_index').getAttribute("img_index")
120
+ gradioApp().dispatchEvent(new Event("image_browser_get_current_img"))
121
+ await image_browser_unlock()
122
+ if (image_browser_debug) console.log("image_browser_get_current_img:end")
123
+ return [
124
+ tab_base_tag,
125
+ img_index,
126
+ page_index,
127
+ filenames,
128
+ turn_page_switch,
129
+ image_gallery
130
+ ]
131
+ }
132
+
133
+ async function image_browser_refresh_current_page_preview() {
134
+ if (image_browser_debug) console.log("image_browser_refresh_current_page_preview:start")
135
+ await image_browser_delay(200)
136
+ const preview_div = gradioApp().querySelector('.preview')
137
+ if (preview_div === null) {
138
+ if (image_browser_debug) console.log("image_browser_refresh_current_page_preview:end")
139
+ return
140
+ }
141
+ const tab_base_tag = image_browser_current_tab()
142
+ const gallery = gradioApp().querySelector(`#${tab_base_tag}_image_browser`)
143
+ const set_btn = gallery.querySelector(".image_browser_set_index")
144
+ const curr_idx = parseInt(set_btn.getAttribute("img_index"))
145
+ // no loading animation, so click immediately
146
+ const gallery_items = gallery.querySelectorAll(".thumbnail-item")
147
+ const curr_image = gallery_items[curr_idx]
148
+ curr_image.click()
149
+ if (image_browser_debug) console.log("image_browser_refresh_current_page_preview:end")
150
+ }
151
+
152
+ async function image_browser_turnpage(tab_base_tag) {
153
+ if (image_browser_debug) console.log("image_browser_turnpage:start")
154
+ while (!image_browser_started) {
155
+ await image_browser_delay(200)
156
+ }
157
+ const gallery = gradioApp().querySelector(`#${tab_base_tag}_image_browser`)
158
+ let clear
159
+ try {
160
+ clear = gallery.querySelector("button[aria-label='Clear']")
161
+ if (clear) {
162
+ clear.click()
163
+ }
164
+ } catch (e) {
165
+ console.error(e)
166
+ }
167
+ if (image_browser_debug) console.log("image_browser_turnpage:end")
168
+ }
169
+
170
+ const image_browser_get_current_img_handler = (del_img_btn) => {
171
+ if (image_browser_debug) console.log("image_browser_get_current_img_handler:start")
172
+ // Prevent delete button spam
173
+ del_img_btn.style.pointerEvents = "auto"
174
+ del_img_btn.style.cursor = "default"
175
+ del_img_btn.style.opacity = "1"
176
+ if (image_browser_debug) console.log("image_browser_get_current_img_handler:end")
177
+ }
178
+
179
+ async function image_browser_select_image(tab_base_tag, img_index, select_image) {
180
+ if (image_browser_debug) console.log("image_browser_select_image:start")
181
+ if (select_image) {
182
+ await image_browser_lock("image_browser_select_image")
183
+ const del_img_btn = gradioApp().getElementById(tab_base_tag + "_image_browser_del_img_btn")
184
+ // Prevent delete button spam
185
+ del_img_btn.style.pointerEvents = "none"
186
+ del_img_btn.style.cursor = "not-allowed"
187
+ del_img_btn.style.opacity = "0.65"
188
+
189
+ const gallery = gradioApp().getElementById(tab_base_tag + "_image_browser_gallery")
190
+ const gallery_items = gallery.querySelectorAll(".thumbnail-item")
191
+ if (img_index >= gallery_items.length || gallery_items.length == 0) {
192
+ const refreshBtn = gradioApp().getElementById(tab_base_tag + "_image_browser_renew_page")
193
+ refreshBtn.dispatchEvent(new Event("click"))
194
+ } else {
195
+ const curr_image = gallery_items[img_index]
196
+ curr_image.click()
197
+ }
198
+ await image_browser_unlock()
199
+
200
+ // Prevent delete button spam
201
+ gradioApp().removeEventListener("image_browser_get_current_img", () => image_browser_get_current_img_handler(del_img_btn))
202
+ gradioApp().addEventListener("image_browser_get_current_img", () => image_browser_get_current_img_handler(del_img_btn))
203
+ }
204
+ if (image_browser_debug) console.log("image_browser_select_image:end")
205
+ }
206
+
207
+ async function image_browser_gototab(tabname) {
208
+ if (image_browser_debug) console.log("image_browser_gototab:start")
209
+ await image_browser_lock("image_browser_gototab")
210
+
211
+ tabNav = gradioApp().querySelector(".tab-nav")
212
+ const tabNavChildren = tabNav.children
213
+ let tabNavButtonNum
214
+ if (typeof tabname === "number") {
215
+ let buttonCnt = 0
216
+ for (let i = 0; i < tabNavChildren.length; i++) {
217
+ if (tabNavChildren[i].tagName === "BUTTON") {
218
+ if (buttonCnt === tabname) {
219
+ tabNavButtonNum = i
220
+ break
221
+ }
222
+ buttonCnt++
223
+ }
224
+ }
225
+ } else {
226
+ for (let i = 0; i < tabNavChildren.length; i++) {
227
+ if (tabNavChildren[i].tagName === "BUTTON" && tabNavChildren[i].textContent.trim() === tabname) {
228
+ tabNavButtonNum = i
229
+ break
230
+ }
231
+ }
232
+ }
233
+ let tabNavButton = tabNavChildren[tabNavButtonNum]
234
+ tabNavButton.click()
235
+
236
+ // Wait for click-action to complete
237
+ const startTime = Date.now()
238
+ // 60 seconds in milliseconds
239
+ const timeout = 60000
240
+
241
+ await image_browser_delay(100)
242
+ while (!tabNavButton.classList.contains("selected")) {
243
+ tabNavButton = tabNavChildren[tabNavButtonNum]
244
+ if (Date.now() - startTime > timeout) {
245
+ throw new Error("image_browser_gototab: 60 seconds have passed")
246
+ }
247
+ await image_browser_delay(200)
248
+ }
249
+
250
+ await image_browser_unlock()
251
+ if (image_browser_debug) console.log("image_browser_gototab:end")
252
+ }
253
+
254
+ async function image_browser_get_image_for_ext(tab_base_tag, image_index) {
255
+ if (image_browser_debug) console.log("image_browser_get_image_for_ext:start")
256
+ const image_browser_image = gradioApp().querySelectorAll(`#${tab_base_tag}_image_browser_gallery .thumbnail-item`)[image_index]
257
+
258
+ const canvas = document.createElement("canvas")
259
+ const image = document.createElement("img")
260
+ image.src = image_browser_image.querySelector("img").src
261
+
262
+ await image.decode()
263
+
264
+ canvas.width = image.width
265
+ canvas.height = image.height
266
+
267
+ canvas.getContext("2d").drawImage(image, 0, 0)
268
+
269
+ if (image_browser_debug) console.log("image_browser_get_image_for_ext:end")
270
+ return canvas.toDataURL()
271
+ }
272
+
273
+ function image_browser_openoutpaint_send(tab_base_tag, image_index, image_browser_prompt, image_browser_neg_prompt, name = "WebUI Resource") {
274
+ if (image_browser_debug) console.log("image_browser_openoutpaint_send:start")
275
+ image_browser_get_image_for_ext(tab_base_tag, image_index)
276
+ .then((dataURL) => {
277
+ // Send to openOutpaint
278
+ openoutpaint_send_image(dataURL, name)
279
+
280
+ // Send prompt to openOutpaint
281
+ const tab = get_uiCurrentTabContent().id
282
+
283
+ const prompt = image_browser_prompt
284
+ const negPrompt = image_browser_neg_prompt
285
+ openoutpaint.frame.contentWindow.postMessage({
286
+ key: openoutpaint.key,
287
+ type: "openoutpaint/set-prompt",
288
+ prompt,
289
+ negPrompt,
290
+ })
291
+
292
+ // Change Tab
293
+ image_browser_gototab("openOutpaint")
294
+ })
295
+ if (image_browser_debug) console.log("image_browser_openoutpaint_send:end")
296
+ }
297
+
298
+ async function image_browser_controlnet_send(toTabNum, tab_base_tag, image_index, controlnetNum, controlnetType) {
299
+ if (image_browser_debug) console.log("image_browser_controlnet_send:start")
300
+ // Logic originally based on github.com/fkunn1326/openpose-editor
301
+ const dataURL = await image_browser_get_image_for_ext(tab_base_tag, image_index)
302
+ const blob = await (await fetch(dataURL)).blob()
303
+ const dt = new DataTransfer()
304
+ dt.items.add(new File([blob], "ImageBrowser.png", { type: blob.type }))
305
+ const list = dt.files
306
+
307
+ await image_browser_gototab(toTabNum)
308
+ const current_tabid = image_browser_webui_current_tab()
309
+ const current_tab = current_tabid.replace("tab_", "")
310
+ const tab_controlnet = gradioApp().getElementById(current_tab + "_controlnet")
311
+ let accordion = tab_controlnet.querySelector("#controlnet > .label-wrap > .icon")
312
+ if (accordion.style.transform.includes("rotate(90deg)")) {
313
+ accordion.click()
314
+ // Wait for click-action to complete
315
+ const startTime = Date.now()
316
+ // 60 seconds in milliseconds
317
+ const timeout = 60000
318
+
319
+ await image_browser_delay(100)
320
+ while (accordion.style.transform.includes("rotate(90deg)")) {
321
+ accordion = tab_controlnet.querySelector("#controlnet > .label-wrap > .icon")
322
+ if (Date.now() - startTime > timeout) {
323
+ throw new Error("image_browser_controlnet_send/accordion: 60 seconds have passed")
324
+ }
325
+ await image_browser_delay(200)
326
+ }
327
+ }
328
+
329
+ let inputImage
330
+ let inputContainer
331
+ if (controlnetType == "single") {
332
+ inputImage = gradioApp().getElementById(current_tab + "_controlnet_ControlNet_input_image")
333
+ } else {
334
+ const tabs = gradioApp().getElementById(current_tab + "_controlnet_tabs")
335
+ const tab_num = (parseInt(controlnetNum) + 1).toString()
336
+ tab_button = tabs.querySelector(".tab-nav > button:nth-child(" + tab_num + ")")
337
+ tab_button.click()
338
+ // Wait for click-action to complete
339
+ const startTime = Date.now()
340
+ // 60 seconds in milliseconds
341
+ const timeout = 60000
342
+
343
+ await image_browser_delay(100)
344
+ while (!tab_button.classList.contains("selected")) {
345
+ tab_button = tabs.querySelector(".tab-nav > button:nth-child(" + tab_num + ")")
346
+ if (Date.now() - startTime > timeout) {
347
+ throw new Error("image_browser_controlnet_send/tabs: 60 seconds have passed")
348
+ }
349
+ await image_browser_delay(200)
350
+ }
351
+ inputImage = gradioApp().getElementById(current_tab + "_controlnet_ControlNet-" + controlnetNum.toString() + "_input_image")
352
+ }
353
+ try {
354
+ inputContainer = inputImage.querySelector('div[data-testid="image"]')
355
+ } catch (e) {}
356
+
357
+ const input = inputContainer.querySelector("input[type='file']")
358
+
359
+ let clear
360
+ try {
361
+ clear = inputContainer.querySelector("button[aria-label='Clear']")
362
+ if (clear) {
363
+ clear.click()
364
+ }
365
+ } catch (e) {
366
+ console.error(e)
367
+ }
368
+
369
+ try {
370
+ // Wait for click-action to complete
371
+ const startTime = Date.now()
372
+ // 60 seconds in milliseconds
373
+ const timeout = 60000
374
+ while (clear) {
375
+ clear = inputContainer.querySelector("button[aria-label='Clear']")
376
+ if (Date.now() - startTime > timeout) {
377
+ throw new Error("image_browser_controlnet_send/clear: 60 seconds have passed")
378
+ }
379
+ await image_browser_delay(200)
380
+ }
381
+ } catch (e) {
382
+ console.error(e)
383
+ }
384
+
385
+ input.value = ""
386
+ input.files = list
387
+ const event = new Event("change", { "bubbles": true, "composed": true })
388
+ input.dispatchEvent(event)
389
+ if (image_browser_debug) console.log("image_browser_controlnet_send:end")
390
+ }
391
+
392
+ function image_browser_controlnet_send_txt2img(tab_base_tag, image_index, controlnetNum, controlnetType) {
393
+ image_browser_controlnet_send(0, tab_base_tag, image_index, controlnetNum, controlnetType)
394
+ }
395
+
396
+ function image_browser_controlnet_send_img2img(tab_base_tag, image_index, controlnetNum, controlnetType) {
397
+ image_browser_controlnet_send(1, tab_base_tag, image_index, controlnetNum, controlnetType)
398
+ }
399
+
400
+ function image_browser_class_add(tab_base_tag) {
401
+ gradioApp().getElementById(tab_base_tag + '_image_browser').classList.add("image_browser_container")
402
+ gradioApp().getElementById(tab_base_tag + '_image_browser_set_index').classList.add("image_browser_set_index")
403
+ gradioApp().getElementById(tab_base_tag + '_image_browser_del_img_btn').classList.add("image_browser_del_img_btn")
404
+ gradioApp().getElementById(tab_base_tag + '_image_browser_gallery').classList.add("image_browser_gallery")
405
+ }
406
+
407
+ function btnClickHandler(tab_base_tag, btn) {
408
+ if (image_browser_debug) console.log("btnClickHandler:start")
409
+ const tabs_box = gradioApp().getElementById("image_browser_tabs_container")
410
+ if (!tabs_box.classList.contains(tab_base_tag)) {
411
+ gradioApp().getElementById(tab_base_tag + "_image_browser_renew_page").click()
412
+ tabs_box.classList.add(tab_base_tag)
413
+ }
414
+ if (image_browser_debug) console.log("btnClickHandler:end")
415
+ }
416
+
417
+ function image_browser_init() {
418
+ if (image_browser_debug) console.log("image_browser_init:start")
419
+ const tab_base_tags = gradioApp().getElementById("image_browser_tab_base_tags_list")
420
+ if (tab_base_tags) {
421
+ const image_browser_tab_base_tags_list = tab_base_tags.querySelector("textarea").value.split(",")
422
+ image_browser_tab_base_tags_list.forEach(function(tab_base_tag) {
423
+ image_browser_class_add(tab_base_tag)
424
+ })
425
+
426
+ const tab_btns = gradioApp().getElementById("image_browser_tabs_container").querySelector("div").querySelectorAll("button")
427
+ tab_btns.forEach(function(btn, i) {
428
+ const tab_base_tag = image_browser_tab_base_tags_list[i]
429
+ btn.setAttribute("tab_base_tag", tab_base_tag)
430
+ btn.removeEventListener('click', () => btnClickHandler(tab_base_tag, btn))
431
+ btn.addEventListener('click', () => btnClickHandler(tab_base_tag, btn))
432
+ })
433
+ //preload
434
+ if (gradioApp().getElementById("image_browser_preload").querySelector("input").checked) {
435
+ setTimeout(function(){tab_btns[0].click()}, 100)
436
+ }
437
+ }
438
+ image_browser_keydown()
439
+
440
+ const image_browser_swipe = gradioApp().getElementById("image_browser_swipe").getElementsByTagName("input")[0]
441
+ if (image_browser_swipe.checked) {
442
+ image_browser_touch()
443
+ }
444
+ if (image_browser_debug) console.log("image_browser_init:end")
445
+ }
446
+
447
+ async function image_browser_wait_for_gallery_btn(tab_base_tag){
448
+ if (image_browser_debug) console.log("image_browser_wait_for_gallery_btn:start")
449
+ await image_browser_delay(100)
450
+ while (!gradioApp().getElementById(image_browser_current_tab() + "_image_browser_gallery").getElementsByClassName("thumbnail-item")) {
451
+ await image_browser_delay(200)
452
+ }
453
+ if (image_browser_debug) console.log("image_browser_wait_for_gallery_btn:end")
454
+ }
455
+
456
+ function image_browser_hijack_console_log() {
457
+ (function () {
458
+ const oldLog = console.log
459
+ console.log = function (message) {
460
+ const formattedTime = new Date().toISOString().slice(0, -5).replace(/[TZ]/g, ' ').trim().replace(/\s+/g, '-').replace(/:/g, '-')
461
+ image_browser_console_log = image_browser_console_log + formattedTime + " " + "image_browser.js: " + message + "\n"
462
+ oldLog.apply(console, arguments)
463
+ }
464
+ })()
465
+ image_browser_debug = true
466
+ }
467
+
468
+ function get_js_logs() {
469
+ log_to_py = image_browser_console_log
470
+ image_browser_console_log = ""
471
+ return log_to_py
472
+ }
473
+
474
+ function isNumeric(str) {
475
+ if (typeof str != "string") return false
476
+ return !isNaN(str) && !isNaN(parseFloat(str))
477
+ }
478
+
479
+ function image_browser_start() {
480
+ if (image_browser_debug) console.log("image_browser_start:start")
481
+ image_browser_init()
482
+ const mutationObserver = new MutationObserver(function(mutationsList) {
483
+ const tab_base_tags = gradioApp().getElementById("image_browser_tab_base_tags_list")
484
+ if (tab_base_tags) {
485
+ const image_browser_tab_base_tags_list = tab_base_tags.querySelector("textarea").value.split(",")
486
+ image_browser_tab_base_tags_list.forEach(function(tab_base_tag) {
487
+ image_browser_class_add(tab_base_tag)
488
+ const tab_gallery_items = gradioApp().querySelectorAll('#' + tab_base_tag + '_image_browser .thumbnail-item')
489
+
490
+ const image_browser_img_info_json = gradioApp().getElementById(tab_base_tag + "_image_browser_img_info").querySelector('[data-testid="textbox"]').value
491
+ const image_browser_img_info = JSON.parse(image_browser_img_info_json)
492
+ const filenames = Object.keys(image_browser_img_info)
493
+
494
+ tab_gallery_items.forEach(function(gallery_item, i) {
495
+ gallery_item.removeEventListener('click', image_browser_click_image, true)
496
+ gallery_item.addEventListener('click', image_browser_click_image, true)
497
+
498
+ const filename = filenames[i]
499
+ try {
500
+ let x = image_browser_img_info[filename].x
501
+ let y = image_browser_img_info[filename].y
502
+ if (isNumeric(x) && isNumeric(y)) {
503
+ gallery_item.title = x + "x" + y
504
+ }
505
+ } catch (e) {}
506
+
507
+ document.onkeyup = async function(e) {
508
+ if (!image_browser_active()) {
509
+ if (image_browser_debug) console.log("image_browser_start:end")
510
+ return
511
+ }
512
+ const current_tab = image_browser_current_tab()
513
+ image_browser_wait_for_gallery_btn(current_tab).then(() => {
514
+ let gallery_btn
515
+ gallery_btn = gradioApp().getElementById(current_tab + "_image_browser_gallery").querySelector(".thumbnail-item .selected")
516
+ gallery_btn = gallery_btn && gallery_btn.length > 0 ? gallery_btn[0] : null
517
+ if (gallery_btn) {
518
+ image_browser_click_image.call(gallery_btn)
519
+ }
520
+ })
521
+ }
522
+ })
523
+
524
+ const cls_btn = gradioApp().getElementById(tab_base_tag + '_image_browser_gallery').querySelector("svg")
525
+ if (cls_btn) {
526
+ cls_btn.removeEventListener('click', () => image_browser_renew_page(tab_base_tag), false)
527
+ cls_btn.addEventListener('click', () => image_browser_renew_page(tab_base_tag), false)
528
+ }
529
+ })
530
+ const debug_level_option = gradioApp().getElementById("image_browser_debug_level_option").querySelector("textarea").value
531
+ if ((debug_level_option == 'javascript' || debug_level_option == 'capture') && !image_browser_debug) {
532
+ image_browser_hijack_console_log()
533
+ }
534
+ }
535
+ })
536
+ mutationObserver.observe(gradioApp(), { childList:true, subtree:true })
537
+ image_browser_started = true
538
+ image_browser_activate_controls()
539
+ if (image_browser_debug) console.log("image_browser_start:end")
540
+ }
541
+
542
+ async function image_browser_activate_controls() {
543
+ if (image_browser_debug) console.log("image_browser_activate_controls:start")
544
+ await image_browser_delay(500)
545
+ container = gradioApp().getElementById("image_browser_tabs_container")
546
+ let controls = container.querySelectorAll('[id*="_control_"]')
547
+ controls.forEach(function(control) {
548
+ control.style.pointerEvents = "auto"
549
+ control.style.cursor = "default"
550
+ control.style.opacity = "1"
551
+ })
552
+ let warnings = container.querySelectorAll('[id*="_warning_box"]')
553
+ warnings.forEach(function(warning) {
554
+ warning.innerHTML = "<p>&nbsp"
555
+ })
556
+ if (image_browser_debug) console.log("image_browser_activate_controls:end")
557
+ }
558
+
559
+ function image_browser_img_show_progress_update() {
560
+ image_browser_img_show_in_progress = false
561
+ }
562
+
563
+ function image_browser_renew_page(tab_base_tag) {
564
+ if (image_browser_debug) console.log("image_browser_renew_page:start")
565
+ gradioApp().getElementById(tab_base_tag + '_image_browser_renew_page').click()
566
+ if (image_browser_debug) console.log("image_browser_renew_page:end")
567
+ }
568
+
569
+ function image_browser_current_tab() {
570
+ if (image_browser_debug) console.log("image_browser_current_tab:start")
571
+ const tabs = gradioApp().getElementById("image_browser_tabs_container").querySelectorAll('[id$="_image_browser_container"]')
572
+ const tab_base_tags = gradioApp().getElementById("image_browser_tab_base_tags_list")
573
+ const image_browser_tab_base_tags_list = tab_base_tags.querySelector("textarea").value.split(",").sort((a, b) => b.length - a.length)
574
+ for (const element of tabs) {
575
+ if (element.style.display === "block") {
576
+ const id = element.id
577
+ const tab_base_tag = image_browser_tab_base_tags_list.find(element => id.startsWith(element)) || null
578
+ if (image_browser_debug) console.log("image_browser_current_tab:end")
579
+ return tab_base_tag
580
+ }
581
+ }
582
+ if (image_browser_debug) console.log("image_browser_current_tab:end")
583
+ }
584
+
585
+ function image_browser_webui_current_tab() {
586
+ if (image_browser_debug) console.log("image_browser_webui_current_tab:start")
587
+ const tabs = gradioApp().getElementById("tabs").querySelectorAll('[id^="tab_"]')
588
+ let id
589
+ for (const element of tabs) {
590
+ if (element.style.display === "block") {
591
+ id = element.id
592
+ break
593
+ }
594
+ }
595
+ if (image_browser_debug) console.log("image_browser_webui_current_tab:end")
596
+ return id
597
+ }
598
+
599
+ function image_browser_active() {
600
+ if (image_browser_debug) console.log("image_browser_active:start")
601
+ const ext_active = gradioApp().getElementById("tab_image_browser")
602
+ if (image_browser_debug) console.log("image_browser_active:end")
603
+ return ext_active && ext_active.style.display !== "none"
604
+ }
605
+
606
+ async function image_browser_delete_key(tab_base_tag) {
607
+ // Wait for img_show to end
608
+ const startTime = Date.now()
609
+ // 60 seconds in milliseconds
610
+ const timeout = 60000
611
+
612
+ await image_browser_delay(100)
613
+ while (image_browser_img_show_in_progress) {
614
+ if (Date.now() - startTime > timeout) {
615
+ throw new Error("image_browser_delete_key: 60 seconds have passed")
616
+ }
617
+ await image_browser_delay(200)
618
+ }
619
+
620
+ const deleteBtn = gradioApp().getElementById(tab_base_tag + "_image_browser_del_img_btn")
621
+ deleteBtn.dispatchEvent(new Event("click"))
622
+ }
623
+
624
+ function image_browser_keydown() {
625
+ if (image_browser_debug) console.log("image_browser_keydown:start")
626
+ gradioApp().addEventListener("keydown", function(event) {
627
+ // If we are not on the Image Browser Extension, dont listen for keypresses
628
+ if (!image_browser_active()) {
629
+ if (image_browser_debug) console.log("image_browser_keydown:end")
630
+ return
631
+ }
632
+
633
+ // If the user is typing in an input field, dont listen for keypresses
634
+ let target
635
+ if (!event.composed) { // We shouldn't get here as the Shadow DOM is always active, but just in case
636
+ target = event.target
637
+ } else {
638
+ target = event.composedPath()[0]
639
+ }
640
+ if (!target || target.nodeName === "INPUT" || target.nodeName === "TEXTAREA") {
641
+ if (image_browser_debug) console.log("image_browser_keydown:end")
642
+ return
643
+ }
644
+
645
+ const tab_base_tag = image_browser_current_tab()
646
+
647
+ // Listens for keypresses 0-5 and updates the corresponding ranking (0 is the last option, None)
648
+ if (event.code >= "Digit0" && event.code <= "Digit5") {
649
+ const selectedValue = event.code.charAt(event.code.length - 1)
650
+ const radioInputs = gradioApp().getElementById(tab_base_tag + "_control_image_browser_ranking").getElementsByTagName("input")
651
+ for (const input of radioInputs) {
652
+ if (input.value === selectedValue || (selectedValue === '0' && input === radioInputs[radioInputs.length - 1])) {
653
+ input.checked = true
654
+ input.dispatchEvent(new Event("change"))
655
+ break
656
+ }
657
+ }
658
+ }
659
+
660
+ const mod_keys = gradioApp().querySelector(`#${tab_base_tag}_image_browser_mod_keys textarea`).value
661
+ let modifiers_pressed = false
662
+ if (mod_keys.indexOf("C") !== -1 && mod_keys.indexOf("S") !== -1) {
663
+ if (event.ctrlKey && event.shiftKey) {
664
+ modifiers_pressed = true
665
+ }
666
+ } else if (mod_keys.indexOf("S") !== -1) {
667
+ if (!event.ctrlKey && event.shiftKey) {
668
+ modifiers_pressed = true
669
+ }
670
+ } else {
671
+ if (event.ctrlKey && !event.shiftKey) {
672
+ modifiers_pressed = true
673
+ }
674
+ }
675
+
676
+ let modifiers_none = false
677
+ if (!event.ctrlKey && !event.shiftKey && !event.altKey && !event.metaKey) {
678
+ modifiers_none = true
679
+ }
680
+
681
+ if (event.code == "KeyF" && modifiers_none) {
682
+ if (tab_base_tag == "image_browser_tab_favorites") {
683
+ if (image_browser_debug) console.log("image_browser_keydown:end")
684
+ return
685
+ }
686
+ const favoriteBtn = gradioApp().getElementById(tab_base_tag + "_image_browser_favorites_btn")
687
+ favoriteBtn.dispatchEvent(new Event("click"))
688
+ }
689
+
690
+ if (event.code == "KeyR" && modifiers_none) {
691
+ const refreshBtn = gradioApp().getElementById(tab_base_tag + "_image_browser_renew_page")
692
+ refreshBtn.dispatchEvent(new Event("click"))
693
+ }
694
+
695
+ if (event.code == "Delete" && modifiers_none) {
696
+ image_browser_delete_key(tab_base_tag)
697
+ }
698
+
699
+ if (event.code == "ArrowLeft" && modifiers_pressed) {
700
+ const prevBtn = gradioApp().getElementById(tab_base_tag + "_control_image_browser_prev_page")
701
+ prevBtn.dispatchEvent(new Event("click"))
702
+ }
703
+
704
+ if (event.code == "ArrowLeft" && modifiers_none) {
705
+ image_browser_img_show_in_progress = true
706
+ const tab_base_tag = image_browser_current_tab()
707
+ const set_btn = gradioApp().querySelector(`#${tab_base_tag}_image_browser .image_browser_set_index`)
708
+ const curr_idx = parseInt(set_btn.getAttribute("img_index"))
709
+ set_btn.setAttribute("img_index", curr_idx - 1)
710
+ image_browser_refresh_current_page_preview()
711
+ }
712
+
713
+ if (event.code == "ArrowRight" && modifiers_pressed) {
714
+ const nextBtn = gradioApp().getElementById(tab_base_tag + "_control_image_browser_next_page")
715
+ nextBtn.dispatchEvent(new Event("click"))
716
+ }
717
+
718
+ if (event.code == "ArrowRight" && modifiers_none) {
719
+ image_browser_img_show_in_progress = true
720
+ const tab_base_tag = image_browser_current_tab()
721
+ const set_btn = gradioApp().querySelector(`#${tab_base_tag}_image_browser .image_browser_set_index`)
722
+ const curr_idx = parseInt(set_btn.getAttribute("img_index"))
723
+ set_btn.setAttribute("img_index", curr_idx + 1)
724
+ image_browser_refresh_current_page_preview()
725
+ }
726
+ })
727
+ if (image_browser_debug) console.log("image_browser_keydown:end")
728
+ }
729
+
730
+ function image_browser_touch() {
731
+ if (image_browser_debug) console.log("image_browser_touch:start")
732
+ let touchStartX = 0
733
+ let touchEndX = 0
734
+ gradioApp().addEventListener("touchstart", function(event) {
735
+ if (!image_browser_active()) {
736
+ if (image_browser_debug) console.log("image_browser_touch:end")
737
+ return
738
+ }
739
+ touchStartX = event.touches[0].clientX;
740
+ })
741
+ gradioApp().addEventListener("touchend", function(event) {
742
+ if (!image_browser_active()) {
743
+ if (image_browser_debug) console.log("image_browser_touch:end")
744
+ return
745
+ }
746
+ touchEndX = event.changedTouches[0].clientX
747
+ const touchDiffX = touchStartX - touchEndX
748
+ if (touchDiffX > 50) {
749
+ const tab_base_tag = image_browser_current_tab()
750
+ const set_btn = gradioApp().querySelector(`#${tab_base_tag}_image_browser .image_browser_set_index`)
751
+ const curr_idx = parseInt(set_btn.getAttribute("img_index"))
752
+ if (curr_idx >= 1) {
753
+ set_btn.setAttribute("img_index", curr_idx - 1)
754
+ image_browser_refresh_current_page_preview()
755
+ }
756
+ } else if (touchDiffX < -50) {
757
+ const tab_base_tag = image_browser_current_tab()
758
+ const gallery = gradioApp().querySelector(`#${tab_base_tag}_image_browser`)
759
+ const gallery_items = gallery.querySelectorAll(".thumbnail-item")
760
+ const thumbnails = gallery_items.length / 2
761
+ const set_btn = gradioApp().querySelector(`#${tab_base_tag}_image_browser .image_browser_set_index`)
762
+ const curr_idx = parseInt(set_btn.getAttribute("img_index"))
763
+ if (curr_idx + 1 < thumbnails) {
764
+ set_btn.setAttribute("img_index", curr_idx + 1)
765
+ image_browser_refresh_current_page_preview()
766
+ }
767
+ }
768
+ })
769
+ if (image_browser_debug) console.log("image_browser_touch:end")
770
+ }
scripts/image_browser.py ADDED
@@ -0,0 +1,1715 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import csv
3
+ import importlib
4
+ import json
5
+ import logging
6
+ import math
7
+ import os
8
+ import platform
9
+ import random
10
+ import re
11
+ import shutil
12
+ import stat
13
+ import subprocess as sp
14
+ import sys
15
+ import tempfile
16
+ import time
17
+ import torch
18
+ import traceback
19
+ import hashlib
20
+ import modules.extras
21
+ import modules.images
22
+ import modules.ui
23
+ from datetime import datetime
24
+ from modules import paths, shared, script_callbacks, scripts, images
25
+ from modules.shared import opts, cmd_opts
26
+ from modules.ui_common import plaintext_to_html
27
+ from modules.ui_components import ToolButton, DropdownMulti
28
+ from PIL import Image, UnidentifiedImageError
29
+ from packaging import version
30
+ from pathlib import Path
31
+ from typing import List, Tuple
32
+ from itertools import chain
33
+ from io import StringIO
34
+
35
+ try:
36
+ from scripts.wib import wib_db
37
+ except ModuleNotFoundError:
38
+ sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "scripts")))
39
+ from wib import wib_db
40
+
41
+ try:
42
+ from send2trash import send2trash
43
+ send2trash_installed = True
44
+ except ImportError:
45
+ print("Image Browser: send2trash is not installed. recycle bin cannot be used.")
46
+ send2trash_installed = False
47
+
48
+ try:
49
+ import ImageReward
50
+ # temporarily deactivated
51
+ # image_reward_installed = True
52
+ image_reward_installed = False
53
+ except ImportError:
54
+ print("Image Browser: ImageReward is not installed, cannot be used.")
55
+ image_reward_installed = False
56
+
57
+ # Force reload wib_db, as it doesn't get reloaded otherwise, if an extension update is started from webui
58
+ importlib.reload(wib_db)
59
+
60
+ yappi_do = False
61
+
62
+ components_list = ["Sort by", "Filename keyword search", "EXIF keyword search", "Ranking Filter", "Aesthestic Score", "Generation Info", "File Name", "File Time", "Open Folder", "Send to buttons", "Copy to directory", "Gallery Controls Bar", "Ranking Bar", "Delete Bar", "Additional Generation Info"]
63
+
64
+ num_of_imgs_per_page = 0
65
+ loads_files_num = 0
66
+ image_ext_list = [".png", ".jpg", ".jpeg", ".bmp", ".gif", ".webp", ".svg"]
67
+ finfo_aes = {}
68
+ finfo_image_reward = {}
69
+ exif_cache = {}
70
+ finfo_exif = {}
71
+ aes_cache = {}
72
+ image_reward_cache = {}
73
+ none_select = "Nothing selected"
74
+ refresh_symbol = '\U0001f504' # 🔄
75
+ up_symbol = '\U000025b2' # ▲
76
+ down_symbol = '\U000025bc' # ▼
77
+ caution_symbol = '\U000026a0' # ⚠
78
+ folder_symbol = '\U0001f4c2' # 📂
79
+ current_depth = 0
80
+ init = True
81
+ copy_move = ["Move", "Copy"]
82
+ copied_moved = ["Moved", "Copied"]
83
+ np = "negative_prompt: "
84
+ openoutpaint = False
85
+ controlnet = False
86
+ js_dummy_return = None
87
+ log_file = os.path.join(scripts.basedir(), "image_browser.log")
88
+ image_reward_model = None
89
+
90
+ def check_image_browser_active_tabs():
91
+ # Check if Maintenance tab has been added to settings in addition to as a mandatory tab. If so, remove.
92
+ if hasattr(opts, "image_browser_active_tabs"):
93
+ active_tabs_no_maint = re.sub(r",\s*Maintenance", "", opts.image_browser_active_tabs)
94
+ if len(active_tabs_no_maint) != len(opts.image_browser_active_tabs):
95
+ shared.opts.__setattr__("image_browser_active_tabs", active_tabs_no_maint)
96
+ shared.opts.save(shared.config_filename)
97
+
98
+ favorite_tab_name = "Favorites"
99
+ default_tab_options = ["txt2img", "img2img", "txt2img-grids", "img2img-grids", "Extras", favorite_tab_name, "Others"]
100
+ check_image_browser_active_tabs()
101
+ tabs_list = [tab.strip() for tab in chain.from_iterable(csv.reader(StringIO(opts.image_browser_active_tabs))) if tab] if hasattr(opts, "image_browser_active_tabs") else default_tab_options
102
+ try:
103
+ if opts.image_browser_enable_maint:
104
+ tabs_list.append("Maintenance") # mandatory tab
105
+ except AttributeError:
106
+ tabs_list.append("Maintenance") # mandatory tab
107
+
108
+ path_maps = {
109
+ "txt2img": opts.outdir_samples or opts.outdir_txt2img_samples,
110
+ "img2img": opts.outdir_samples or opts.outdir_img2img_samples,
111
+ "txt2img-grids": opts.outdir_grids or opts.outdir_txt2img_grids,
112
+ "img2img-grids": opts.outdir_grids or opts.outdir_img2img_grids,
113
+ "Extras": opts.outdir_samples or opts.outdir_extras_samples,
114
+ favorite_tab_name: opts.outdir_save
115
+ }
116
+
117
+ class ImageBrowserTab():
118
+
119
+ seen_base_tags = set()
120
+
121
+ def __init__(self, name: str):
122
+ self.name: str = os.path.basename(name) if os.path.isdir(name) else name
123
+ self.path: str = os.path.realpath(path_maps.get(name, name))
124
+ self.base_tag: str = f"image_browser_tab_{self.get_unique_base_tag(self.remove_invalid_html_tag_chars(self.name).lower())}"
125
+
126
+ def remove_invalid_html_tag_chars(self, tag: str) -> str:
127
+ # Removes any character that is not a letter, a digit, a hyphen, or an underscore
128
+ removed = re.sub(r'[^a-zA-Z0-9\-_]', '', tag)
129
+ return removed
130
+
131
+ def get_unique_base_tag(self, base_tag: str) -> str:
132
+ counter = 1
133
+ while base_tag in self.seen_base_tags:
134
+ match = re.search(r'_(\d+)$', base_tag)
135
+ if match:
136
+ counter = int(match.group(1)) + 1
137
+ base_tag = re.sub(r'_(\d+)$', f"_{counter}", base_tag)
138
+ else:
139
+ base_tag = f"{base_tag}_{counter}"
140
+ counter += 1
141
+ self.seen_base_tags.add(base_tag)
142
+ return base_tag
143
+
144
+ def __str__(self):
145
+ return f"Name: {self.name} / Path: {self.path} / Base tag: {self.base_tag} / Seen base tags: {self.seen_base_tags}"
146
+
147
+ tabs_list = [ImageBrowserTab(tab) for tab in tabs_list]
148
+
149
+ debug_level_types = ["none", "warning log", "debug log", "javascript log", "capture logs to file"]
150
+
151
+ debug_levels_list = []
152
+ for i in range(len(debug_level_types)):
153
+ level = debug_level_types[i].split(" ")[0]
154
+ text = str(i) + " - " + debug_level_types[i]
155
+ debug_levels_list.append((level, text))
156
+
157
+ def debug_levels(arg_value=None, arg_level=None, arg_text=None):
158
+ if arg_value is not None:
159
+ return arg_value, debug_levels_list[arg_value]
160
+ elif arg_level is not None:
161
+ for i, (level, text) in enumerate(debug_levels_list):
162
+ if level == arg_level:
163
+ return i, debug_levels_list[i]
164
+ elif arg_text is not None:
165
+ for i, (level, text) in enumerate(debug_levels_list):
166
+ if text == arg_text:
167
+ return i, debug_levels_list[i]
168
+
169
+ # Logging
170
+ logger = None
171
+ def restart_debug(parameter):
172
+ global logger
173
+ logger = logging.getLogger(__name__)
174
+ logger.disabled = False
175
+ logger_mode = logging.ERROR
176
+ level_value = 0
177
+ capture_level_value = 99
178
+ if hasattr(opts, "image_browser_debug_level"):
179
+ warning_level_value, (warning_level, warning_level_text) = debug_levels(arg_level="warning")
180
+ debug_level_value, (debug_level, debug_level_text) = debug_levels(arg_level="debug")
181
+ capture_level_value, (capture_level, capture_level_text) = debug_levels(arg_level="capture")
182
+ level_value, (level, level_text) = debug_levels(arg_text=opts.image_browser_debug_level)
183
+ if level_value >= debug_level_value:
184
+ logger_mode = logging.DEBUG
185
+ elif level_value >= warning_level_value:
186
+ logger_mode = logging.WARNING
187
+ logger.setLevel(logger_mode)
188
+ if (logger.hasHandlers()):
189
+ logger.handlers.clear()
190
+ console_handler = logging.StreamHandler()
191
+ console_handler.setLevel(logger_mode)
192
+ formatter = logging.Formatter(f'%(asctime)s image_browser.py: %(message)s', datefmt='%Y-%m-%d-%H:%M:%S')
193
+ console_handler.setFormatter(formatter)
194
+ logger.addHandler(console_handler)
195
+ if level_value >= capture_level_value:
196
+ try:
197
+ os.unlink(log_file)
198
+ except FileNotFoundError:
199
+ pass
200
+ file_handler = logging.FileHandler(log_file)
201
+ file_handler.setLevel(logger_mode)
202
+ file_handler.setFormatter(formatter)
203
+ logger.addHandler(file_handler)
204
+ logger.warning(f"debug_level: {level_value}")
205
+ # Debug logging
206
+ if logger.getEffectiveLevel() == logging.DEBUG:
207
+ if parameter != "startup":
208
+ logging.disable(logging.NOTSET)
209
+
210
+ logger.debug(f"{sys.executable} {sys.version}")
211
+ logger.debug(f"{platform.system()} {platform.version()}")
212
+ try:
213
+ git = os.environ.get('GIT', "git")
214
+ webui_commit_hash = os.popen(f"{git} rev-parse HEAD").read().strip()
215
+ sm_hashes = os.popen(f"{git} submodule").read()
216
+ sm_hashes_lines = sm_hashes.splitlines()
217
+ image_browser_commit_hash = f"image_browser_commit_hash not found: {sm_hashes}"
218
+ for sm_hashes_line in sm_hashes_lines:
219
+ if "images-browser" in sm_hashes_line.lower():
220
+ image_browser_commit_hash = sm_hashes_line[1:41]
221
+ break
222
+ except Exception as e:
223
+ webui_commit_hash = e
224
+ image_browser_commit_hash = e
225
+ logger.debug(f"Webui {webui_commit_hash}")
226
+ logger.debug(f"Image Browser {image_browser_commit_hash}")
227
+ logger.debug(f"Gradio {gr.__version__}")
228
+ logger.debug(f"{paths.script_path}")
229
+ with open(cmd_opts.ui_config_file, "r") as f:
230
+ logger.debug(f.read())
231
+ with open(cmd_opts.ui_settings_file, "r") as f:
232
+ logger.debug(f.read())
233
+ logger.debug(os.path.realpath(__file__))
234
+ logger.debug([str(tab) for tab in tabs_list])
235
+ maint_last_msg = "Debug restarted"
236
+
237
+ return parameter, maint_last_msg
238
+
239
+ restart_debug("startup")
240
+
241
+ def delete_recycle(filename):
242
+ if opts.image_browser_delete_recycle and send2trash_installed:
243
+ send2trash(filename)
244
+ else:
245
+ file = Path(filename)
246
+ file.unlink()
247
+ return
248
+
249
+ def img_path_subdirs_get(img_path):
250
+ subdirs = []
251
+ subdirs.append(none_select)
252
+ for item in os.listdir(img_path):
253
+ item_path = os.path.join(img_path, item)
254
+ if os.path.isdir(item_path):
255
+ subdirs.append(item_path)
256
+ return gr.update(choices=subdirs)
257
+
258
+ def img_path_add_remove(img_dir, path_recorder, add_remove, img_path_depth):
259
+ img_dir = os.path.realpath(img_dir)
260
+ if add_remove == "add" or (add_remove == "remove" and img_dir in path_recorder):
261
+ if add_remove == "add":
262
+ path_recorder[img_dir] = {
263
+ "depth": int(img_path_depth),
264
+ "path_display": f"{img_dir} [{int(img_path_depth)}]"
265
+ }
266
+ wib_db.update_path_recorder(img_dir, path_recorder[img_dir]["depth"], path_recorder[img_dir]["path_display"])
267
+ else:
268
+ del path_recorder[img_dir]
269
+ wib_db.delete_path_recorder(img_dir)
270
+ path_recorder_formatted = [value.get("path_display") for key, value in path_recorder.items()]
271
+ path_recorder_formatted = sorted(path_recorder_formatted, key=lambda x: natural_keys(x.lower()))
272
+
273
+ if add_remove == "remove":
274
+ selected = path_recorder[list(path_recorder.keys())[0]]["path_display"]
275
+ else:
276
+ selected = path_recorder[img_dir]["path_display"]
277
+ return path_recorder, gr.update(choices=path_recorder_formatted, value=selected)
278
+
279
+ def sort_order_flip(turn_page_switch, sort_order):
280
+ if sort_order == up_symbol:
281
+ sort_order = down_symbol
282
+ else:
283
+ sort_order = up_symbol
284
+ return 1, -turn_page_switch, sort_order
285
+
286
+ def read_path_recorder():
287
+ path_recorder = wib_db.load_path_recorder()
288
+ path_recorder_formatted = [value.get("path_display") for key, value in path_recorder.items()]
289
+ path_recorder_formatted = sorted(path_recorder_formatted, key=lambda x: natural_keys(x.lower()))
290
+ path_recorder_unformatted = list(path_recorder.keys())
291
+ path_recorder_unformatted = sorted(path_recorder_unformatted, key=lambda x: natural_keys(x.lower()))
292
+
293
+ return path_recorder, path_recorder_formatted, path_recorder_unformatted
294
+
295
+ def pure_path(path):
296
+ if path == []:
297
+ return path, 0
298
+ match = re.search(r" \[(\d+)\]$", path)
299
+ if match:
300
+ path = path[:match.start()]
301
+ depth = int(match.group(1))
302
+ else:
303
+ depth = 0
304
+ path = os.path.realpath(path)
305
+ return path, depth
306
+
307
+ def browser2path(img_path_browser):
308
+ img_path, _ = pure_path(img_path_browser)
309
+ return img_path
310
+
311
+ def totxt(file):
312
+ base, _ = os.path.splitext(file)
313
+ file_txt = base + '.txt'
314
+
315
+ return file_txt
316
+
317
+ def tab_select():
318
+ path_recorder, path_recorder_formatted, path_recorder_unformatted = read_path_recorder()
319
+ return path_recorder, gr.update(choices=path_recorder_unformatted)
320
+
321
+ def js_logs_output(js_log):
322
+ logger.debug(f"js_log: {js_log}")
323
+ return js_log
324
+
325
+ def ranking_filter_settings(page_index, turn_page_switch, ranking_filter):
326
+ if ranking_filter == "Min-max":
327
+ interactive = True
328
+ else:
329
+ interactive = False
330
+ page_index = 1
331
+ turn_page_switch = -turn_page_switch
332
+ return page_index, turn_page_switch, gr.update(interactive=interactive), gr.update(interactive=interactive)
333
+
334
+ def reduplicative_file_move(src, dst):
335
+ def same_name_file(basename, path):
336
+ name, ext = os.path.splitext(basename)
337
+ f_list = os.listdir(path)
338
+ max_num = 0
339
+ for f in f_list:
340
+ if len(f) <= len(basename):
341
+ continue
342
+ f_ext = f[-len(ext):] if len(ext) > 0 else ""
343
+ if f[:len(name)] == name and f_ext == ext:
344
+ if f[len(name)] == "(" and f[-len(ext)-1] == ")":
345
+ number = f[len(name)+1:-len(ext)-1]
346
+ if number.isdigit():
347
+ if int(number) > max_num:
348
+ max_num = int(number)
349
+ return f"{name}({max_num + 1}){ext}"
350
+ name = os.path.basename(src)
351
+ save_name = os.path.join(dst, name)
352
+ src_txt_exists = False
353
+ if opts.image_browser_txt_files:
354
+ src_txt = totxt(src)
355
+ if os.path.exists(src_txt):
356
+ src_txt_exists = True
357
+ if not os.path.exists(save_name):
358
+ if opts.image_browser_copy_image:
359
+ shutil.copy2(src, dst)
360
+ if opts.image_browser_txt_files and src_txt_exists:
361
+ shutil.copy2(src_txt, dst)
362
+ else:
363
+ shutil.move(src, dst)
364
+ if opts.image_browser_txt_files and src_txt_exists:
365
+ shutil.move(src_txt, dst)
366
+ else:
367
+ name = same_name_file(name, dst)
368
+ if opts.image_browser_copy_image:
369
+ shutil.copy2(src, os.path.join(dst, name))
370
+ if opts.image_browser_txt_files and src_txt_exists:
371
+ shutil.copy2(src_txt, totxt(os.path.join(dst, name)))
372
+ else:
373
+ shutil.move(src, os.path.join(dst, name))
374
+ if opts.image_browser_txt_files and src_txt_exists:
375
+ shutil.move(src_txt, totxt(os.path.join(dst, name)))
376
+
377
+ def save_image(file_name, filenames, page_index, turn_page_switch, dest_path):
378
+ if file_name is not None and os.path.exists(file_name):
379
+ reduplicative_file_move(file_name, dest_path)
380
+ message = f"<div style='color:#999'>{copied_moved[opts.image_browser_copy_image]} to {dest_path}</div>"
381
+ if not opts.image_browser_copy_image:
382
+ # Force page refresh with checking filenames
383
+ filenames = []
384
+ turn_page_switch = -turn_page_switch
385
+ else:
386
+ message = "<div style='color:#999'>Image not found (may have been already moved)</div>"
387
+
388
+ return message, filenames, page_index, turn_page_switch
389
+
390
+ def delete_image(tab_base_tag_box, delete_num, name, filenames, image_index, visible_num, delete_confirm, turn_page_switch, image_page_list):
391
+ logger.debug("delete_image")
392
+ refresh = False
393
+ delete_num = int(delete_num)
394
+ image_index = int(image_index)
395
+ visible_num = int(visible_num)
396
+ image_page_list = json.loads(image_page_list)
397
+ new_file_list = []
398
+ new_image_page_list = []
399
+ if name == "":
400
+ refresh = True
401
+ else:
402
+ try:
403
+ index_files = list(filenames).index(name)
404
+
405
+ index_on_page = image_page_list.index(name)
406
+ except ValueError as e:
407
+ print(traceback.format_exc(), file=sys.stderr)
408
+ # Something went wrong, force a page refresh
409
+ refresh = True
410
+ if not refresh:
411
+ if not delete_confirm:
412
+ delete_num = min(visible_num - index_on_page, delete_num)
413
+ new_file_list = filenames[:index_files] + filenames[index_files + delete_num:]
414
+ new_image_page_list = image_page_list[:index_on_page] + image_page_list[index_on_page + delete_num:]
415
+
416
+ for i in range(index_files, index_files + delete_num):
417
+ if os.path.exists(filenames[i]):
418
+ if opts.image_browser_delete_message:
419
+ print(f"Deleting file {filenames[i]}")
420
+ delete_recycle(filenames[i])
421
+ visible_num -= 1
422
+ if opts.image_browser_txt_files:
423
+ txt_file = totxt(filenames[i])
424
+ if os.path.exists(txt_file):
425
+ delete_recycle(txt_file)
426
+ else:
427
+ print(f"File does not exist {filenames[i]}")
428
+ # If we reach this point (which we shouldn't), things are messed up, better force a page refresh
429
+ refresh = True
430
+
431
+ if refresh:
432
+ turn_page_switch = -turn_page_switch
433
+ select_image = False
434
+ else:
435
+ select_image = True
436
+
437
+ return new_file_list, 1, turn_page_switch, visible_num, new_image_page_list, select_image, json.dumps(new_image_page_list)
438
+
439
+ def traverse_all_files(curr_path, image_list, tab_base_tag_box, img_path_depth) -> List[Tuple[str, os.stat_result, str, int]]:
440
+ global current_depth
441
+ logger.debug(f"curr_path: {curr_path}")
442
+ if curr_path == "":
443
+ return image_list
444
+ f_list = [(os.path.join(curr_path, entry.name), entry.stat()) for entry in os.scandir(curr_path)]
445
+ for f_info in f_list:
446
+ fname, fstat = f_info
447
+ if os.path.splitext(fname)[1] in image_ext_list:
448
+ image_list.append(f_info)
449
+ elif stat.S_ISDIR(fstat.st_mode):
450
+ if (opts.image_browser_with_subdirs and tab_base_tag_box != "image_browser_tab_others") or (tab_base_tag_box == "image_browser_tab_others" and img_path_depth != 0 and (current_depth < img_path_depth or img_path_depth < 0)):
451
+ current_depth = current_depth + 1
452
+ image_list = traverse_all_files(fname, image_list, tab_base_tag_box, img_path_depth)
453
+ current_depth = current_depth - 1
454
+ return image_list
455
+
456
+ def cache_exif(fileinfos):
457
+ global finfo_exif, exif_cache, finfo_aes, aes_cache, finfo_image_reward, image_reward_cache
458
+
459
+ if yappi_do:
460
+ import yappi
461
+ import pandas as pd
462
+ yappi.set_clock_type("wall")
463
+ yappi.start()
464
+
465
+ cache_exif_start = time.time()
466
+ new_exif = 0
467
+ new_aes = 0
468
+ conn, cursor = wib_db.transaction_begin()
469
+ for fi_info in fileinfos:
470
+ if any(fi_info[0].endswith(ext) for ext in image_ext_list):
471
+ found_exif = False
472
+ found_aes = False
473
+ if fi_info[0] in exif_cache:
474
+ finfo_exif[fi_info[0]] = exif_cache[fi_info[0]]
475
+ found_exif = True
476
+ if fi_info[0] in aes_cache:
477
+ finfo_aes[fi_info[0]] = aes_cache[fi_info[0]]
478
+ found_aes = True
479
+ if fi_info[0] in image_reward_cache:
480
+ finfo_image_reward[fi_info[0]] = image_reward_cache[fi_info[0]]
481
+ if not found_exif or not found_aes:
482
+ finfo_exif[fi_info[0]] = "0"
483
+ exif_cache[fi_info[0]] = "0"
484
+ finfo_aes[fi_info[0]] = "0"
485
+ aes_cache[fi_info[0]] = "0"
486
+ try:
487
+ image = Image.open(fi_info[0])
488
+ (_, allExif, allExif_html) = modules.extras.run_pnginfo(image)
489
+ image.close()
490
+ except SyntaxError:
491
+ allExif = False
492
+ logger.warning(f"Extension and content don't match: {fi_info[0]}")
493
+ except UnidentifiedImageError as e:
494
+ allExif = False
495
+ logger.warning(f"UnidentifiedImageError: {e}")
496
+ except Image.DecompressionBombError as e:
497
+ allExif = False
498
+ logger.warning(f"DecompressionBombError: {e}: {fi_info[0]}")
499
+ except PermissionError as e:
500
+ allExif = False
501
+ logger.warning(f"PermissionError: {e}: {fi_info[0]}")
502
+ except FileNotFoundError as e:
503
+ allExif = False
504
+ logger.warning(f"FileNotFoundError: {e}: {fi_info[0]}")
505
+ except OSError as e:
506
+ if e.errno == 22:
507
+ logger.warning(f"Caught OSError with error code 22: {fi_info[0]}")
508
+ else:
509
+ raise
510
+ if allExif:
511
+ finfo_exif[fi_info[0]] = allExif
512
+ exif_cache[fi_info[0]] = allExif
513
+ wib_db.update_exif_data(conn, fi_info[0], allExif)
514
+ new_exif = new_exif + 1
515
+
516
+ m = re.search("(?:aesthetic_score:|Score:) (\d+.\d+)", allExif)
517
+ if m:
518
+ aes_value = m.group(1)
519
+ else:
520
+ aes_value = "0"
521
+ finfo_aes[fi_info[0]] = aes_value
522
+ aes_cache[fi_info[0]] = aes_value
523
+ wib_db.update_exif_data_by_key(conn, fi_info[0], "aesthetic_score", aes_value)
524
+ new_aes = new_aes + 1
525
+ else:
526
+ try:
527
+ filename = os.path.splitext(fi_info[0])[0] + ".txt"
528
+ geninfo = ""
529
+ with open(filename) as f:
530
+ for line in f:
531
+ geninfo += line
532
+ finfo_exif[fi_info[0]] = geninfo
533
+ exif_cache[fi_info[0]] = geninfo
534
+ wib_db.update_exif_data_by_key(conn, fi_info[0], geninfo)
535
+ new_exif = new_exif + 1
536
+
537
+ m = re.search("(?:aesthetic_score:|Score:) (\d+.\d+)", geninfo)
538
+ if m:
539
+ aes_value = m.group(1)
540
+ else:
541
+ aes_value = "0"
542
+ finfo_aes[fi_info[0]] = aes_value
543
+ aes_cache[fi_info[0]] = aes_value
544
+ wib_db.update_exif_data_by_key(conn, fi_info[0], "aesthetic_score", aes_value)
545
+ new_aes = new_aes + 1
546
+ except Exception:
547
+ logger.warning(f"cache_exif: No EXIF in image or txt file for {fi_info[0]}")
548
+ # Saved with defaults to not scan it again next time
549
+ finfo_exif[fi_info[0]] = "0"
550
+ exif_cache[fi_info[0]] = "0"
551
+ allExif = "0"
552
+ wib_db.update_exif_data(conn, fi_info[0], allExif)
553
+ new_exif = new_exif + 1
554
+
555
+ aes_value = "0"
556
+ finfo_aes[fi_info[0]] = aes_value
557
+ aes_cache[fi_info[0]] = aes_value
558
+ wib_db.update_exif_data_by_key(conn, fi_info[0], "aesthetic_score", aes_value)
559
+ new_aes = new_aes + 1
560
+ wib_db.transaction_end(conn, cursor)
561
+
562
+ if yappi_do:
563
+ yappi.stop()
564
+ pd.set_option('display.float_format', lambda x: '%.6f' % x)
565
+ yappi_stats = yappi.get_func_stats().strip_dirs()
566
+ data = [(s.name, s.ncall, s.tsub, s.ttot, s.ttot/s.ncall) for s in yappi_stats]
567
+ df = pd.DataFrame(data, columns=['name', 'ncall', 'tsub', 'ttot', 'tavg'])
568
+ print(df.to_string(index=False))
569
+ yappi.get_thread_stats().print_all()
570
+
571
+ cache_exif_end = time.time()
572
+ logger.debug(f"cache_exif: {new_exif}/{len(fileinfos)} cache_aes: {new_aes}/{len(fileinfos)} {round(cache_exif_end - cache_exif_start, 1)} seconds")
573
+
574
+ def exif_rebuild(maint_wait):
575
+ global finfo_exif, exif_cache, finfo_aes, aes_cache, finfo_image_reward, image_reward_cache
576
+ if opts.image_browser_scan_exif:
577
+ logger.debug("Rebuild start")
578
+ exif_dirs = wib_db.get_exif_dirs()
579
+ finfo_aes = {}
580
+ finfo_image_reward = {}
581
+ exif_cache = {}
582
+ finfo_exif = {}
583
+ aes_cache = {}
584
+ image_reward_cache = {}
585
+ for key, value in exif_dirs.items():
586
+ if os.path.exists(key):
587
+ print(f"Rebuilding {key}")
588
+ fileinfos = traverse_all_files(key, [], "", 0)
589
+ cache_exif(fileinfos)
590
+ logger.debug("Rebuild end")
591
+ maint_last_msg = "Rebuild finished"
592
+ else:
593
+ maint_last_msg = "Exif cache not enabled in settings"
594
+
595
+ return maint_wait, maint_last_msg
596
+
597
+ def exif_delete_0(maint_wait):
598
+ global finfo_exif, exif_cache, finfo_aes, aes_cache
599
+ if opts.image_browser_scan_exif:
600
+ conn, cursor = wib_db.transaction_begin()
601
+ wib_db.delete_exif_0(cursor)
602
+ wib_db.transaction_end(conn, cursor)
603
+ finfo_aes = {}
604
+ finfo_exif = {}
605
+ exif_cache = wib_db.load_exif_data(exif_cache)
606
+ aes_cache = wib_db.load_aes_data(aes_cache)
607
+ maint_last_msg = "Delete finished"
608
+ else:
609
+ maint_last_msg = "Exif cache not enabled in settings"
610
+
611
+ return maint_wait, maint_last_msg
612
+
613
+ def exif_update_dirs(maint_update_dirs_from, maint_update_dirs_to, maint_wait):
614
+ global exif_cache, aes_cache, image_reward_cache
615
+ if maint_update_dirs_from == "":
616
+ maint_last_msg = "From is empty"
617
+ elif maint_update_dirs_to == "":
618
+ maint_last_msg = "To is empty"
619
+ else:
620
+ maint_update_dirs_from = os.path.realpath(maint_update_dirs_from)
621
+ maint_update_dirs_to = os.path.realpath(maint_update_dirs_to)
622
+ rows = 0
623
+ conn, cursor = wib_db.transaction_begin()
624
+ wib_db.update_path_recorder_mult(cursor, maint_update_dirs_from, maint_update_dirs_to)
625
+ rows = rows + cursor.rowcount
626
+ wib_db.update_exif_data_mult(cursor, maint_update_dirs_from, maint_update_dirs_to)
627
+ rows = rows + cursor.rowcount
628
+ wib_db.update_ranking_mult(cursor, maint_update_dirs_from, maint_update_dirs_to)
629
+ rows = rows + cursor.rowcount
630
+ wib_db.transaction_end(conn, cursor)
631
+ if rows == 0:
632
+ maint_last_msg = "No rows updated"
633
+ else:
634
+ maint_last_msg = f"{rows} rows updated. Please reload UI!"
635
+
636
+ return maint_wait, maint_last_msg
637
+
638
+ def reapply_ranking(path_recorder, maint_wait):
639
+ dirs = {}
640
+
641
+ for tab in tabs_list:
642
+ if os.path.exists(tab.path):
643
+ dirs[tab.path] = tab.path
644
+
645
+ for key in path_recorder:
646
+ if os.path.exists(key):
647
+ dirs[key] = key
648
+
649
+ conn, cursor = wib_db.transaction_begin()
650
+
651
+ # Traverse all known dirs, check if missing rankings are due to moved files
652
+ for key in dirs.keys():
653
+ fileinfos = traverse_all_files(key, [], "", 0)
654
+ for (file, _) in fileinfos:
655
+ # Is there a ranking for this full filepath
656
+ ranking_by_file = wib_db.get_ranking_by_file(cursor, file)
657
+ if ranking_by_file is None:
658
+ name = os.path.basename(file)
659
+ (ranking_by_name, alternate_hash) = wib_db.get_ranking_by_name(cursor, name)
660
+ # Is there a ranking only for the filename
661
+ if ranking_by_name is not None:
662
+ hash = wib_db.get_hash(file)
663
+ (alternate_file, alternate_ranking) = ranking_by_name
664
+ if alternate_ranking is not None:
665
+ (alternate_hash,) = alternate_hash
666
+ # Does the found filename's file have no hash or the same hash?
667
+ if alternate_hash is None or hash == alternate_hash:
668
+ if os.path.exists(alternate_file):
669
+ # Insert ranking as a copy of the found filename's ranking
670
+ wib_db.insert_ranking(cursor, file, alternate_ranking, hash)
671
+ else:
672
+ # Replace ranking of the found filename
673
+ wib_db.replace_ranking(cursor, file, alternate_file, hash)
674
+
675
+ wib_db.transaction_end(conn, cursor)
676
+ maint_last_msg = "Rankings reapplied"
677
+
678
+ return maint_wait, maint_last_msg
679
+
680
+ def atof(text):
681
+ try:
682
+ retval = float(text)
683
+ except ValueError:
684
+ retval = text
685
+ return retval
686
+
687
+ def natural_keys(text):
688
+ '''
689
+ alist.sort(key=natural_keys) sorts in human order
690
+ http://nedbatchelder.com/blog/200712/human_sorting.html
691
+ (See Toothy's implementation in the comments)
692
+ float regex comes from https://stackoverflow.com/a/12643073/190597
693
+ '''
694
+ return [ atof(c) for c in re.split(r'[+-]?([0-9]+(?:[.][0-9]*)?|[.][0-9]+)', text) ]
695
+
696
+ def open_folder(path):
697
+ if os.path.exists(path):
698
+ # Code from ui_common.py
699
+ if not shared.cmd_opts.hide_ui_dir_config:
700
+ if platform.system() == "Windows":
701
+ os.startfile(path)
702
+ elif platform.system() == "Darwin":
703
+ sp.Popen(["open", path])
704
+ elif "microsoft-standard-WSL2" in platform.uname().release:
705
+ sp.Popen(["wsl-open", path])
706
+ else:
707
+ sp.Popen(["xdg-open", path])
708
+
709
+ def check_ext(ext):
710
+ found = False
711
+ scripts_list = scripts.list_scripts("scripts", ".py")
712
+ for scriptfile in scripts_list:
713
+ if ext in scriptfile.basedir.lower():
714
+ found = True
715
+ break
716
+ return found
717
+
718
+ def exif_search(needle, haystack, use_regex, case_sensitive):
719
+ found = False
720
+ if use_regex:
721
+ if case_sensitive:
722
+ pattern = re.compile(needle, re.DOTALL)
723
+ else:
724
+ pattern = re.compile(needle, re.DOTALL | re.IGNORECASE)
725
+ if pattern.search(haystack) is not None:
726
+ found = True
727
+ else:
728
+ if not case_sensitive:
729
+ haystack = haystack.lower()
730
+ needle = needle.lower()
731
+ if needle in haystack:
732
+ found = True
733
+ return found
734
+
735
+ def get_all_images(dir_name, sort_by, sort_order, keyword, tab_base_tag_box, img_path_depth, ranking_filter, ranking_filter_min, ranking_filter_max, aes_filter_min, aes_filter_max, score_type, exif_keyword, negative_prompt_search, use_regex, case_sensitive):
736
+ global current_depth
737
+ logger.debug("get_all_images")
738
+ current_depth = 0
739
+ fileinfos = traverse_all_files(dir_name, [], tab_base_tag_box, img_path_depth)
740
+ keyword = keyword.strip(" ")
741
+
742
+ if opts.image_browser_scan_exif:
743
+ cache_exif(fileinfos)
744
+
745
+ if len(keyword) != 0:
746
+ fileinfos = [x for x in fileinfos if keyword.lower() in x[0].lower()]
747
+ filenames = [finfo[0] for finfo in fileinfos]
748
+
749
+ if opts.image_browser_scan_exif:
750
+ conn, cursor = wib_db.transaction_begin()
751
+ if len(exif_keyword) != 0:
752
+ if use_regex:
753
+ regex_error = False
754
+ try:
755
+ test_re = re.compile(exif_keyword, re.DOTALL)
756
+ except re.error as e:
757
+ regex_error = True
758
+ print(f"Regex error: {e}")
759
+ if (use_regex and not regex_error) or not use_regex:
760
+ if negative_prompt_search == "Yes":
761
+ fileinfos = [x for x in fileinfos if exif_search(exif_keyword, finfo_exif[x[0]], use_regex, case_sensitive)]
762
+ else:
763
+ result = []
764
+ for file_info in fileinfos:
765
+ file_name = file_info[0]
766
+ file_exif = finfo_exif[file_name]
767
+ file_exif_lc = file_exif.lower()
768
+ start_index = file_exif_lc.find(np)
769
+ end_index = file_exif.find("\n", start_index)
770
+ if negative_prompt_search == "Only":
771
+ start_index = start_index + len(np)
772
+ sub_string = file_exif[start_index:end_index].strip()
773
+ if exif_search(exif_keyword, sub_string, use_regex, case_sensitive):
774
+ result.append(file_info)
775
+ else:
776
+ sub_string = file_exif[start_index:end_index].strip()
777
+ file_exif = file_exif.replace(sub_string, "")
778
+
779
+ if exif_search(exif_keyword, file_exif, use_regex, case_sensitive):
780
+ result.append(file_info)
781
+ fileinfos = result
782
+ filenames = [finfo[0] for finfo in fileinfos]
783
+ wib_db.fill_work_files(cursor, fileinfos)
784
+ if len(aes_filter_min) != 0 or len(aes_filter_max) != 0:
785
+ try:
786
+ aes_filter_min_num = float(aes_filter_min)
787
+ except ValueError:
788
+ aes_filter_min_num = sys.float_info.min
789
+ try:
790
+ aes_filter_max_num = float(aes_filter_max)
791
+ except ValueError:
792
+ aes_filter_max_num = sys.float_info.max
793
+
794
+ fileinfos = wib_db.filter_aes(cursor, fileinfos, aes_filter_min_num, aes_filter_max_num, score_type)
795
+ filenames = [finfo[0] for finfo in fileinfos]
796
+ if ranking_filter != "All":
797
+ ranking_filter_min_num = 1
798
+ ranking_filter_max_num = 5
799
+ if ranking_filter == "Min-max":
800
+ try:
801
+ ranking_filter_min_num = int(ranking_filter_min)
802
+ except ValueError:
803
+ ranking_filter_min_num = 0
804
+ try:
805
+ ranking_filter_max_num = int(ranking_filter_max)
806
+ except ValueError:
807
+ ranking_filter_max_num = 0
808
+ if ranking_filter_min_num < 1:
809
+ ranking_filter_min_num = 1
810
+ if ranking_filter_max_num < 1 or ranking_filter_max_num > 5:
811
+ ranking_filter_max_num = 5
812
+
813
+ fileinfos = wib_db.filter_ranking(cursor, fileinfos, ranking_filter, ranking_filter_min_num, ranking_filter_max_num)
814
+ filenames = [finfo[0] for finfo in fileinfos]
815
+
816
+ wib_db.transaction_end(conn, cursor)
817
+
818
+ if sort_by == "date":
819
+ if sort_order == up_symbol:
820
+ fileinfos = sorted(fileinfos, key=lambda x: x[1].st_mtime)
821
+ else:
822
+ fileinfos = sorted(fileinfos, key=lambda x: -x[1].st_mtime)
823
+ filenames = [finfo[0] for finfo in fileinfos]
824
+ elif sort_by == "path name":
825
+ if sort_order == up_symbol:
826
+ fileinfos = sorted(fileinfos)
827
+ else:
828
+ fileinfos = sorted(fileinfos, reverse=True)
829
+ filenames = [finfo[0] for finfo in fileinfos]
830
+ elif sort_by == "random":
831
+ random.shuffle(fileinfos)
832
+ filenames = [finfo[0] for finfo in fileinfos]
833
+ elif sort_by == "ranking":
834
+ finfo_ranked = {}
835
+ for fi_info in fileinfos:
836
+ finfo_ranked[fi_info[0]], _ = get_ranking(fi_info[0])
837
+ if sort_order == up_symbol:
838
+ fileinfos = dict(sorted(finfo_ranked.items(), key=lambda x: (x[1], x[0])))
839
+ else:
840
+ fileinfos = dict(reversed(sorted(finfo_ranked.items(), key=lambda x: (x[1], x[0]))))
841
+ filenames = [finfo for finfo in fileinfos]
842
+ else:
843
+ sort_values = {}
844
+ exif_info = dict(finfo_exif)
845
+ if exif_info:
846
+ for k, v in exif_info.items():
847
+ match = re.search(r'(?<='+ sort_by + ":" ').*?(?=(,|$))', v, flags=re.DOTALL|re.IGNORECASE)
848
+ if match:
849
+ sort_values[k] = match.group().strip()
850
+ else:
851
+ sort_values[k] = "0"
852
+ if sort_by == "aesthetic_score" or sort_by == "ImageRewardScore" or sort_by == "cfg scale":
853
+ sort_float = True
854
+ else:
855
+ sort_float = False
856
+
857
+ if sort_order == down_symbol:
858
+ if sort_float:
859
+ fileinfos = [x for x in fileinfos if sort_values[x[0]] != "0"]
860
+ fileinfos.sort(key=lambda x: float(sort_values[x[0]]), reverse=True)
861
+ fileinfos = dict(fileinfos)
862
+ else:
863
+ fileinfos = dict(reversed(sorted(fileinfos, key=lambda x: natural_keys(sort_values[x[0]]))))
864
+ else:
865
+ if sort_float:
866
+ fileinfos = [x for x in fileinfos if sort_values[x[0]] != "0"]
867
+ fileinfos.sort(key=lambda x: float(sort_values[x[0]]))
868
+ fileinfos = dict(fileinfos)
869
+ else:
870
+ fileinfos = dict(sorted(fileinfos, key=lambda x: natural_keys(sort_values[x[0]])))
871
+ filenames = [finfo for finfo in fileinfos]
872
+ else:
873
+ filenames = [finfo for finfo in fileinfos]
874
+ return filenames
875
+
876
+ def get_image_thumbnail(image_list):
877
+ logger.debug("get_image_thumbnail")
878
+ optimized_cache = os.path.join(tempfile.gettempdir(),"optimized")
879
+ os.makedirs(optimized_cache,exist_ok=True)
880
+ thumbnail_list = []
881
+ for image_path in image_list:
882
+ image_path_hash = hashlib.md5(image_path.encode("utf-8")).hexdigest()
883
+ cache_image_path = os.path.join(optimized_cache, image_path_hash + ".jpg")
884
+ if os.path.isfile(cache_image_path):
885
+ thumbnail_list.append(cache_image_path)
886
+ else:
887
+ try:
888
+ image = Image.open(image_path)
889
+ except OSError:
890
+ # If PIL cannot open the image, use the original path
891
+ thumbnail_list.append(image_path)
892
+ continue
893
+ width, height = image.size
894
+ left = (width - min(width, height)) / 2
895
+ top = (height - min(width, height)) / 2
896
+ right = (width + min(width, height)) / 2
897
+ bottom = (height + min(width, height)) / 2
898
+ thumbnail = image.crop((left, top, right, bottom))
899
+ thumbnail.thumbnail((opts.image_browser_thumbnail_size, opts.image_browser_thumbnail_size))
900
+ if thumbnail.mode != "RGB":
901
+ thumbnail = thumbnail.convert("RGB")
902
+ try:
903
+ thumbnail.save(cache_image_path, "JPEG")
904
+ thumbnail_list.append(cache_image_path)
905
+ except FileNotFoundError:
906
+ # Cannot save cache, use PIL object
907
+ thumbnail_list.append(thumbnail)
908
+ return thumbnail_list
909
+
910
+ def set_tooltip_info(image_list):
911
+ image_browser_img_info = {}
912
+ conn, cursor = wib_db.transaction_begin()
913
+ for filename in image_list:
914
+ x, y = wib_db.select_x_y(cursor, filename)
915
+ image_browser_img_info[filename] = {"x": x, "y": y}
916
+ wib_db.transaction_end(conn, cursor)
917
+ image_browser_img_info_json = json.dumps(image_browser_img_info)
918
+ return image_browser_img_info_json
919
+
920
+ def get_image_page(img_path, page_index, filenames, keyword, sort_by, sort_order, tab_base_tag_box, img_path_depth, ranking_filter, ranking_filter_min, ranking_filter_max, aes_filter_min, aes_filter_max, score_type, exif_keyword, negative_prompt_search, use_regex, case_sensitive, image_reward_button):
921
+ logger.debug("get_image_page")
922
+ if img_path == "":
923
+ return [], page_index, [], "", "", "", 0, "", None, "", "[]", image_reward_button
924
+
925
+ # Set temp_dir from webui settings, so gradio uses it
926
+ if shared.opts.temp_dir != "":
927
+ tempfile.tempdir = shared.opts.temp_dir
928
+
929
+ img_path, _ = pure_path(img_path)
930
+ filenames = get_all_images(img_path, sort_by, sort_order, keyword, tab_base_tag_box, img_path_depth, ranking_filter, ranking_filter_min, ranking_filter_max, aes_filter_min, aes_filter_max, score_type, exif_keyword, negative_prompt_search, use_regex, case_sensitive)
931
+ page_index = int(page_index)
932
+ length = len(filenames)
933
+ max_page_index = math.ceil(length / num_of_imgs_per_page)
934
+ page_index = max_page_index if page_index == -1 else page_index
935
+ page_index = 1 if page_index < 1 else page_index
936
+ page_index = max_page_index if page_index > max_page_index else page_index
937
+ idx_frm = (page_index - 1) * num_of_imgs_per_page
938
+ image_list = filenames[idx_frm:idx_frm + num_of_imgs_per_page]
939
+
940
+ if opts.image_browser_scan_exif and opts.image_browser_img_tooltips:
941
+ image_browser_img_info = set_tooltip_info(image_list)
942
+ else:
943
+ image_browser_img_info = "[]"
944
+
945
+ if opts.image_browser_use_thumbnail:
946
+ thumbnail_list = get_image_thumbnail(image_list)
947
+ else:
948
+ thumbnail_list = image_list
949
+
950
+ visible_num = num_of_imgs_per_page if idx_frm + num_of_imgs_per_page < length else length % num_of_imgs_per_page
951
+ visible_num = num_of_imgs_per_page if visible_num == 0 else visible_num
952
+
953
+ load_info = "<div style='color:#999' align='center'>"
954
+ load_info += f"{length} images in this directory, divided into {int((length + 1) // num_of_imgs_per_page + 1)} pages"
955
+ load_info += "</div>"
956
+
957
+ return filenames, gr.update(value=page_index, label=f"Page Index ({page_index}/{max_page_index})"), thumbnail_list, "", "", "", visible_num, load_info, None, json.dumps(image_list), image_browser_img_info, gr.update(visible=True)
958
+
959
+ def get_current_file(tab_base_tag_box, num, page_index, filenames):
960
+ file = filenames[int(num) + int((page_index - 1) * num_of_imgs_per_page)]
961
+ return file
962
+
963
+ def show_image_info(tab_base_tag_box, num, page_index, filenames, turn_page_switch, image_gallery):
964
+ logger.debug(f"show_image_info: tab_base_tag_box, num, page_index, len(filenames), num_of_imgs_per_page: {tab_base_tag_box}, {num}, {page_index}, {len(filenames)}, {num_of_imgs_per_page}")
965
+ if len(filenames) == 0:
966
+ # This should only happen if webui was stopped and started again and the user clicks on one of the still displayed images.
967
+ # The state with the filenames will be empty then. In that case we return None to prevent further errors and force a page refresh.
968
+ turn_page_switch = -turn_page_switch
969
+ file = None
970
+ tm = None
971
+ info = ""
972
+ else:
973
+ file_num = int(num) + int(
974
+ (page_index - 1) * num_of_imgs_per_page)
975
+ if file_num >= len(filenames):
976
+ # Last image to the right is deleted, page refresh
977
+ turn_page_switch = -turn_page_switch
978
+ file = None
979
+ tm = None
980
+ info = ""
981
+ else:
982
+ file = filenames[file_num]
983
+ tm = "<div style='color:#999' align='right'>" + time.strftime("%Y-%m-%d %H:%M:%S",time.localtime(os.path.getmtime(file))) + "</div>"
984
+ try:
985
+ with Image.open(file) as image:
986
+ _, geninfo, info = modules.extras.run_pnginfo(image)
987
+ except UnidentifiedImageError as e:
988
+ info = ""
989
+ logger.warning(f"UnidentifiedImageError: {e}")
990
+ if opts.image_browser_use_thumbnail:
991
+ image_gallery = [image['name'] for image in image_gallery]
992
+ image_gallery[int(num)] = filenames[file_num]
993
+ return file, tm, num, file, turn_page_switch, info, image_gallery
994
+ else:
995
+ return file, tm, num, file, turn_page_switch, info
996
+
997
+ def change_dir(img_dir, path_recorder, load_switch, img_path_browser, img_path_depth, img_path):
998
+ warning = None
999
+ img_path, _ = pure_path(img_path)
1000
+ img_path_depth_org = img_path_depth
1001
+ if img_dir == none_select:
1002
+ return warning, gr.update(visible=False), img_path_browser, path_recorder, load_switch, img_path, img_path_depth
1003
+ else:
1004
+ img_dir, img_path_depth = pure_path(img_dir)
1005
+ if warning is None:
1006
+ try:
1007
+ if os.path.exists(img_dir):
1008
+ try:
1009
+ f = os.listdir(img_dir)
1010
+ except:
1011
+ warning = f"'{img_dir} is not a directory"
1012
+ else:
1013
+ warning = "The directory does not exist"
1014
+ except:
1015
+ warning = "The format of the directory is incorrect"
1016
+ if warning is None:
1017
+ return "", gr.update(visible=True), img_path_browser, path_recorder, img_dir, img_dir, img_path_depth
1018
+ else:
1019
+ return warning, gr.update(visible=False), img_path_browser, path_recorder, load_switch, img_path, img_path_depth_org
1020
+
1021
+ def update_move_text_one(btn):
1022
+ btn_text = " ".join(btn.split()[1:])
1023
+ return f"{copy_move[opts.image_browser_copy_image]} {btn_text}"
1024
+
1025
+ def update_move_text(favorites_btn, to_dir_btn):
1026
+ return update_move_text_one(favorites_btn), update_move_text_one(to_dir_btn)
1027
+
1028
+ def get_ranking(filename):
1029
+ ranking_value = wib_db.select_ranking(filename)
1030
+ return ranking_value, None
1031
+
1032
+ def img_file_name_changed(img_file_name, favorites_btn, to_dir_btn):
1033
+ ranking_current, ranking = get_ranking(img_file_name)
1034
+ favorites_btn, to_dir_btn = update_move_text(favorites_btn, to_dir_btn)
1035
+
1036
+ return ranking_current, ranking, "", favorites_btn, to_dir_btn
1037
+
1038
+ def update_exif(img_file_name, key, value):
1039
+ image = Image.open(img_file_name)
1040
+ geninfo, items = images.read_info_from_image(image)
1041
+ if geninfo is not None:
1042
+ if f"{key}: " in geninfo:
1043
+ if value == "None":
1044
+ geninfo = re.sub(f', {key}: \d+(\.\d+)*', '', geninfo)
1045
+ else:
1046
+ geninfo = re.sub(f'{key}: \d+(\.\d+)*', f'{key}: {value}', geninfo)
1047
+ else:
1048
+ geninfo = f'{geninfo}, {key}: {value}'
1049
+
1050
+ original_time = os.path.getmtime(img_file_name)
1051
+ images.save_image(image, os.path.dirname(img_file_name), "", extension=os.path.splitext(img_file_name)[1][1:], info=geninfo, forced_filename=os.path.splitext(os.path.basename(img_file_name))[0], save_to_dirs=False)
1052
+ os.utime(img_file_name, (original_time, original_time))
1053
+ return geninfo
1054
+
1055
+ def update_ranking(img_file_name, ranking_current, ranking, img_file_info):
1056
+ # ranking = None is different than ranking = "None"! None means no radio button selected. "None" means radio button called "None" selected.
1057
+ if ranking is None:
1058
+ return ranking_current, None, img_file_info
1059
+
1060
+ saved_ranking, _ = get_ranking(img_file_name)
1061
+ if saved_ranking != ranking:
1062
+ wib_db.update_ranking(img_file_name, ranking)
1063
+ if opts.image_browser_ranking_pnginfo and any(img_file_name.endswith(ext) for ext in image_ext_list):
1064
+ img_file_info = update_exif(img_file_name, "Ranking", ranking)
1065
+ return ranking, None, img_file_info
1066
+
1067
+ def generate_image_reward(filenames, turn_page_switch, aes_filter_min, aes_filter_max):
1068
+ global image_reward_model
1069
+ if image_reward_model is None:
1070
+ image_reward_model = ImageReward.load("ImageReward-v1.0")
1071
+ conn, cursor = wib_db.transaction_begin()
1072
+ for filename in filenames:
1073
+ saved_image_reward_score, saved_image_reward_prompt = wib_db.select_image_reward_score(cursor, filename)
1074
+ if saved_image_reward_score is None and saved_image_reward_prompt is not None:
1075
+ with torch.no_grad():
1076
+ image_reward_score = image_reward_model.score(saved_image_reward_prompt, filename)
1077
+ image_reward_score = f"{image_reward_score:.2f}"
1078
+ try:
1079
+ logger.warning(f"Generated ImageRewardScore: {image_reward_score} for {filename}")
1080
+ except UnicodeEncodeError:
1081
+ pass
1082
+ wib_db.update_image_reward_score(cursor, filename, image_reward_score)
1083
+ if any(filename.endswith(ext) for ext in image_ext_list):
1084
+ img_file_info = update_exif(filename, "ImageRewardScore", image_reward_score)
1085
+ wib_db.transaction_end(conn, cursor)
1086
+ return -turn_page_switch, aes_filter_min, aes_filter_max
1087
+
1088
+ def create_tab(tab: ImageBrowserTab, current_gr_tab: gr.Tab):
1089
+ global init, exif_cache, aes_cache, image_reward_cache, openoutpaint, controlnet, js_dummy_return
1090
+ dir_name = None
1091
+ others_dir = False
1092
+ maint = False
1093
+ standard_ui = True
1094
+ path_recorder = {}
1095
+ path_recorder_formatted = []
1096
+ path_recorder_unformatted = []
1097
+
1098
+ if init:
1099
+ db_version = wib_db.check()
1100
+ logger.debug(f"db_version: {db_version}")
1101
+ exif_cache = wib_db.load_exif_data(exif_cache)
1102
+ aes_cache = wib_db.load_exif_data_by_key(aes_cache, "aesthetic_score", "Score")
1103
+ image_reward_cache = wib_db.load_exif_data_by_key(image_reward_cache, "ImageRewardScore", "ImageRewardScore")
1104
+ init = False
1105
+
1106
+ path_recorder, path_recorder_formatted, path_recorder_unformatted = read_path_recorder()
1107
+ openoutpaint = check_ext("openoutpaint")
1108
+ controlnet = check_ext("controlnet")
1109
+
1110
+ if tab.name == "Others":
1111
+ others_dir = True
1112
+ standard_ui = False
1113
+ elif tab.name == "Maintenance":
1114
+ maint = True
1115
+ standard_ui = False
1116
+ else:
1117
+ dir_name = tab.path
1118
+
1119
+ if standard_ui:
1120
+ dir_name = str(Path(dir_name))
1121
+ if not os.path.exists(dir_name):
1122
+ os.makedirs(dir_name)
1123
+
1124
+ with gr.Row():
1125
+ path_recorder = gr.State(path_recorder)
1126
+ with gr.Column(scale=10):
1127
+ warning_box = gr.HTML("<p>&nbsp", elem_id=f"{tab.base_tag}_image_browser_warning_box")
1128
+ with gr.Column(scale=5, visible=(tab.name==favorite_tab_name)):
1129
+ gr.HTML(f"<p>Favorites path from settings: {opts.outdir_save}")
1130
+
1131
+ with gr.Row(visible=others_dir):
1132
+ with gr.Column(scale=10):
1133
+ img_path = gr.Textbox(dir_name, label="Images directory", placeholder="Input images directory", interactive=others_dir)
1134
+ with gr.Column(scale=1):
1135
+ img_path_depth = gr.Number(value="0", label="Sub directory depth")
1136
+ with gr.Column(scale=1):
1137
+ img_path_save_button = gr.Button(value="Add to / replace in saved directories")
1138
+
1139
+ with gr.Row(visible=others_dir):
1140
+ with gr.Column(scale=10):
1141
+ img_path_browser = gr.Dropdown(choices=path_recorder_formatted, label="Saved directories")
1142
+ with gr.Column(scale=1):
1143
+ img_path_remove_button = gr.Button(value="Remove from saved directories")
1144
+
1145
+ with gr.Row(visible=others_dir):
1146
+ with gr.Column(scale=10):
1147
+ img_path_subdirs = gr.Dropdown(choices=[none_select], value=none_select, label="Sub directories", interactive=True, elem_id=f"{tab.base_tag}_img_path_subdirs")
1148
+ with gr.Column(scale=1):
1149
+ img_path_subdirs_button = gr.Button(value="Get sub directories")
1150
+
1151
+ with gr.Row(visible=standard_ui, elem_id=f"{tab.base_tag}_image_browser") as main_panel:
1152
+ with gr.Column():
1153
+ with gr.Row():
1154
+ with gr.Column(scale=2):
1155
+ with gr.Row(elem_id=f"{tab.base_tag}_image_browser_gallery_controls") as gallery_controls_panel:
1156
+ with gr.Column(scale=2, min_width=20):
1157
+ first_page = gr.Button("First Page", elem_id=f"{tab.base_tag}_control_image_browser_first_page")
1158
+ with gr.Column(scale=2, min_width=20):
1159
+ prev_page = gr.Button("Prev Page", elem_id=f"{tab.base_tag}_control_image_browser_prev_page")
1160
+ with gr.Column(scale=2, min_width=20):
1161
+ page_index = gr.Number(value=1, label="Page Index", elem_id=f"{tab.base_tag}_control_image_browser_page_index")
1162
+ with gr.Column(scale=1, min_width=20):
1163
+ refresh_index_button = ToolButton(value=refresh_symbol, elem_id=f"{tab.base_tag}_control_image_browser_refresh_index")
1164
+ with gr.Column(scale=2, min_width=20):
1165
+ next_page = gr.Button("Next Page", elem_id=f"{tab.base_tag}_control_image_browser_next_page")
1166
+ with gr.Column(scale=2, min_width=20):
1167
+ end_page = gr.Button("End Page", elem_id=f"{tab.base_tag}_control_image_browser_end_page")
1168
+ with gr.Row(visible=False) as ranking_panel:
1169
+ with gr.Column(scale=1, min_width=20):
1170
+ ranking_current = gr.Textbox(value="None", label="Current ranking", interactive=False)
1171
+ with gr.Column(scale=4, min_width=20):
1172
+ ranking = gr.Radio(choices=["1", "2", "3", "4", "5", "None"], label="Set ranking to", elem_id=f"{tab.base_tag}_control_image_browser_ranking", interactive=True)
1173
+ with gr.Row():
1174
+ image_gallery = gr.Gallery(show_label=False, elem_id=f"{tab.base_tag}_image_browser_gallery").style(grid=opts.image_browser_page_columns)
1175
+ with gr.Row() as delete_panel:
1176
+ with gr.Column(scale=1):
1177
+ delete_num = gr.Number(value=1, interactive=True, label="delete next", elem_id=f"{tab.base_tag}_image_browser_del_num")
1178
+ delete_confirm = gr.Checkbox(value=False, label="also delete off-screen images")
1179
+ with gr.Column(scale=3):
1180
+ delete = gr.Button('Delete', elem_id=f"{tab.base_tag}_image_browser_del_img_btn")
1181
+ with gr.Row() as info_add_panel:
1182
+ with gr.Accordion("Additional Generation Info", open=False):
1183
+ img_file_info_add = gr.HTML()
1184
+
1185
+ with gr.Column(scale=1):
1186
+ with gr.Row() as sort_panel:
1187
+ sort_by = gr.Dropdown(value="date", choices=["path name", "date", "aesthetic_score", "ImageRewardScore", "random", "cfg scale", "steps", "seed", "sampler", "size", "model", "model hash", "ranking"], label="Sort by")
1188
+ sort_order = ToolButton(value=down_symbol)
1189
+ with gr.Row() as filename_search_panel:
1190
+ filename_keyword_search = gr.Textbox(value="", label="Filename keyword search")
1191
+ with gr.Box() as exif_search_panel:
1192
+ with gr.Row():
1193
+ exif_keyword_search = gr.Textbox(value="", label="EXIF keyword search")
1194
+ negative_prompt_search = gr.Radio(value="No", choices=["No", "Yes", "Only"], label="Search negative prompt", interactive=True)
1195
+ with gr.Row():
1196
+ case_sensitive = gr.Checkbox(value=False, label="case sensitive")
1197
+ use_regex = gr.Checkbox(value=False, label=r"regex - e.g. ^(?!.*Hires).*$")
1198
+ with gr.Box() as ranking_filter_panel:
1199
+ with gr.Row():
1200
+ ranking_filter = gr.Radio(value="All", choices=["All", "1", "2", "3", "4", "5", "None", "Min-max"], label="Ranking filter", interactive=True)
1201
+ with gr.Row():
1202
+ with gr.Column(scale=2, min_width=20):
1203
+ ranking_filter_min = gr.Textbox(value="1", label="Minimum ranking", interactive=False)
1204
+ with gr.Column(scale=2, min_width=20):
1205
+ ranking_filter_max = gr.Textbox(value="5", label="Maximum ranking", interactive=False)
1206
+ with gr.Column(scale=4, min_width=20):
1207
+ gr.Textbox(value="Choose Min-max to activate these controls", label="", interactive=False)
1208
+ with gr.Box() as aesthetic_score_filter_panel:
1209
+ with gr.Row():
1210
+ with gr.Column(scale=4, min_width=20):
1211
+ score_type = gr.Dropdown(value=opts.image_browser_scoring_type, choices=["aesthetic_score", "ImageReward Score"], label="Scoring type", interactive=True)
1212
+ with gr.Column(scale=2, min_width=20):
1213
+ image_reward_button = gr.Button(value="Generate ImageReward Scores for all images", interactive=image_reward_installed, visible=False)
1214
+ with gr.Row():
1215
+ aes_filter_min = gr.Textbox(value="", label="Minimum score")
1216
+ aes_filter_max = gr.Textbox(value="", label="Maximum score")
1217
+ with gr.Row() as generation_info_panel:
1218
+ img_file_info = gr.Textbox(label="Generation Info", interactive=False, lines=6,elem_id=f"{tab.base_tag}_image_browser_file_info")
1219
+ with gr.Row() as filename_panel:
1220
+ img_file_name = gr.Textbox(value="", label="File Name", interactive=False)
1221
+ with gr.Row() as filetime_panel:
1222
+ img_file_time= gr.HTML()
1223
+ with gr.Row() as open_folder_panel:
1224
+ open_folder_button = gr.Button(folder_symbol, visible=standard_ui or others_dir)
1225
+ gr.HTML("&nbsp")
1226
+ gr.HTML("&nbsp")
1227
+ gr.HTML("&nbsp")
1228
+ with gr.Row(elem_id=f"{tab.base_tag}_image_browser_button_panel", visible=False) as button_panel:
1229
+ with gr.Column():
1230
+ with gr.Row():
1231
+ if tab.name == favorite_tab_name:
1232
+ favorites_btn_show = False
1233
+ else:
1234
+ favorites_btn_show = True
1235
+ favorites_btn = gr.Button(f'{copy_move[opts.image_browser_copy_image]} to favorites', elem_id=f"{tab.base_tag}_image_browser_favorites_btn", visible=favorites_btn_show)
1236
+ try:
1237
+ send_to_buttons = modules.generation_parameters_copypaste.create_buttons(["txt2img", "img2img", "inpaint", "extras"])
1238
+ except:
1239
+ pass
1240
+ sendto_openoutpaint = gr.Button("Send to openOutpaint", elem_id=f"{tab.base_tag}_image_browser_openoutpaint_btn", visible=openoutpaint)
1241
+ with gr.Row(visible=controlnet):
1242
+ sendto_controlnet_txt2img = gr.Button("Send to txt2img ControlNet", visible=controlnet)
1243
+ sendto_controlnet_img2img = gr.Button("Send to img2img ControlNet", visible=controlnet)
1244
+ controlnet_max = opts.data.get("control_net_max_models_num", 1)
1245
+ sendto_controlnet_num = gr.Dropdown([str(i) for i in range(controlnet_max)], label="ControlNet number", value="0", interactive=True, visible=(controlnet and controlnet_max > 1))
1246
+ if controlnet_max is None:
1247
+ sendto_controlnet_type = gr.Textbox(value="none", visible=False)
1248
+ elif controlnet_max == 1:
1249
+ sendto_controlnet_type = gr.Textbox(value="single", visible=False)
1250
+ else:
1251
+ sendto_controlnet_type = gr.Textbox(value="multiple", visible=False)
1252
+ with gr.Row(elem_id=f"{tab.base_tag}_image_browser_to_dir_panel", visible=False) as to_dir_panel:
1253
+ with gr.Box():
1254
+ with gr.Row():
1255
+ to_dir_path = gr.Textbox(label="Directory path")
1256
+ with gr.Row():
1257
+ to_dir_saved = gr.Dropdown(choices=path_recorder_unformatted, label="Saved directories")
1258
+ with gr.Row():
1259
+ to_dir_btn = gr.Button(f'{copy_move[opts.image_browser_copy_image]} to directory', elem_id=f"{tab.base_tag}_image_browser_to_dir_btn")
1260
+
1261
+ with gr.Row():
1262
+ collected_warning = gr.HTML()
1263
+
1264
+ with gr.Row(visible=False):
1265
+ renew_page = gr.Button("Renew Page", elem_id=f"{tab.base_tag}_image_browser_renew_page")
1266
+ visible_img_num = gr.Number()
1267
+ tab_base_tag_box = gr.Textbox(tab.base_tag)
1268
+ image_index = gr.Textbox(value=-1, elem_id=f"{tab.base_tag}_image_browser_image_index")
1269
+ set_index = gr.Button('set_index', elem_id=f"{tab.base_tag}_image_browser_set_index")
1270
+ filenames = gr.State([])
1271
+ hidden = gr.Image(type="pil", elem_id=f"{tab.base_tag}_image_browser_hidden_image")
1272
+ image_page_list = gr.Textbox(elem_id=f"{tab.base_tag}_image_browser_image_page_list")
1273
+ info1 = gr.Textbox()
1274
+ info2 = gr.Textbox()
1275
+ load_switch = gr.Textbox(value="load_switch", label="load_switch")
1276
+ to_dir_load_switch = gr.Textbox(value="to dir load_switch", label="to_dir_load_switch")
1277
+ turn_page_switch = gr.Number(value=1, label="turn_page_switch")
1278
+ select_image = gr.Number(value=1)
1279
+ img_path_add = gr.Textbox(value="add")
1280
+ img_path_remove = gr.Textbox(value="remove")
1281
+ favorites_path = gr.Textbox(value=opts.outdir_save)
1282
+ mod_keys = ""
1283
+ if opts.image_browser_mod_ctrl_shift:
1284
+ mod_keys = f"{mod_keys}CS"
1285
+ elif opts.image_browser_mod_shift:
1286
+ mod_keys = f"{mod_keys}S"
1287
+ image_browser_mod_keys = gr.Textbox(value=mod_keys, elem_id=f"{tab.base_tag}_image_browser_mod_keys")
1288
+ image_browser_prompt = gr.Textbox(elem_id=f"{tab.base_tag}_image_browser_prompt")
1289
+ image_browser_neg_prompt = gr.Textbox(elem_id=f"{tab.base_tag}_image_browser_neg_prompt")
1290
+ js_logs = gr.Textbox()
1291
+ image_browser_img_info = gr.Textbox(value="[]", elem_id=f"{tab.base_tag}_image_browser_img_info")
1292
+
1293
+ # Maintenance tab
1294
+ with gr.Row(visible=maint):
1295
+ with gr.Column(scale=4):
1296
+ gr.HTML(f"{caution_symbol} Caution: You should only use these options if you know what you are doing. {caution_symbol}")
1297
+ with gr.Column(scale=3):
1298
+ maint_wait = gr.HTML("Status:")
1299
+ with gr.Column(scale=7):
1300
+ gr.HTML("&nbsp")
1301
+ with gr.Row(visible=maint):
1302
+ maint_last_msg = gr.Textbox(label="Last message", interactive=False)
1303
+ with gr.Row(visible=maint):
1304
+ with gr.Column(scale=1):
1305
+ maint_exif_rebuild = gr.Button(value="Rebuild exif cache")
1306
+ with gr.Column(scale=1):
1307
+ maint_exif_delete_0 = gr.Button(value="Delete 0-entries from exif cache")
1308
+ with gr.Column(scale=10):
1309
+ gr.HTML(visible=False)
1310
+ with gr.Row(visible=maint):
1311
+ with gr.Column(scale=1):
1312
+ maint_update_dirs = gr.Button(value="Update directory names in database")
1313
+ with gr.Column(scale=10):
1314
+ maint_update_dirs_from = gr.Textbox(label="From (full path)")
1315
+ with gr.Column(scale=10):
1316
+ maint_update_dirs_to = gr.Textbox(label="to (full path)")
1317
+ with gr.Row(visible=maint):
1318
+ with gr.Column(scale=1):
1319
+ maint_reapply_ranking = gr.Button(value="Reapply ranking after moving files")
1320
+ with gr.Column(scale=10):
1321
+ gr.HTML(visible=False)
1322
+ with gr.Row(visible=maint):
1323
+ with gr.Column(scale=1):
1324
+ maint_restart_debug = gr.Button(value="Restart debug")
1325
+ with gr.Column(scale=10):
1326
+ gr.HTML(visible=False)
1327
+ with gr.Row(visible=maint):
1328
+ with gr.Column(scale=1):
1329
+ maint_get_js_logs = gr.Button(value="Get javascript logs")
1330
+ with gr.Column(scale=10):
1331
+ maint_show_logs = gr.Textbox(label="Javascript logs", lines=10, interactive=False)
1332
+ with gr.Row(visible=False):
1333
+ with gr.Column(scale=1):
1334
+ maint_rebuild_ranking = gr.Button(value="Rebuild ranking from exif info")
1335
+ with gr.Column(scale=10):
1336
+ gr.HTML(visible=False)
1337
+
1338
+ # Hide components based on opts.image_browser_hidden_components
1339
+ hidden_component_map = {
1340
+ "Sort by": sort_panel,
1341
+ "Filename keyword search": filename_search_panel,
1342
+ "EXIF keyword search": exif_search_panel,
1343
+ "Ranking Filter": ranking_filter_panel,
1344
+ "Aesthestic Score": aesthetic_score_filter_panel,
1345
+ "Generation Info": generation_info_panel,
1346
+ "File Name": filename_panel,
1347
+ "File Time": filetime_panel,
1348
+ "Open Folder": open_folder_panel,
1349
+ "Send to buttons": button_panel,
1350
+ "Copy to directory": to_dir_panel,
1351
+ "Gallery Controls Bar": gallery_controls_panel,
1352
+ "Ranking Bar": ranking_panel,
1353
+ "Delete Bar": delete_panel,
1354
+ "Additional Generation Info": info_add_panel
1355
+ }
1356
+
1357
+ if set(hidden_component_map.keys()) != set(components_list):
1358
+ logger.warning(f"Invalid items present in either hidden_component_map or components_list. Make sure when adding new components they are added to both.")
1359
+
1360
+ override_hidden = set()
1361
+ if hasattr(opts, "image_browser_hidden_components"):
1362
+ for item in opts.image_browser_hidden_components:
1363
+ hidden_component_map[item].visible = False
1364
+ override_hidden.add(hidden_component_map[item])
1365
+
1366
+ change_dir_outputs = [warning_box, main_panel, img_path_browser, path_recorder, load_switch, img_path, img_path_depth]
1367
+ img_path.submit(change_dir, inputs=[img_path, path_recorder, load_switch, img_path_browser, img_path_depth, img_path], outputs=change_dir_outputs, show_progress=opts.image_browser_show_progress)
1368
+ img_path_browser.change(change_dir, inputs=[img_path_browser, path_recorder, load_switch, img_path_browser, img_path_depth, img_path], outputs=change_dir_outputs, show_progress=opts.image_browser_show_progress)
1369
+ # img_path_browser.change(browser2path, inputs=[img_path_browser], outputs=[img_path])
1370
+ to_dir_saved.change(change_dir, inputs=[to_dir_saved, path_recorder, to_dir_load_switch, to_dir_saved, img_path_depth, to_dir_path], outputs=[warning_box, main_panel, to_dir_saved, path_recorder, to_dir_load_switch, to_dir_path, img_path_depth], show_progress=opts.image_browser_show_progress)
1371
+
1372
+ #delete
1373
+ delete.click(
1374
+ fn=delete_image,
1375
+ inputs=[tab_base_tag_box, delete_num, img_file_name, filenames, image_index, visible_img_num, delete_confirm, turn_page_switch, image_page_list],
1376
+ outputs=[filenames, delete_num, turn_page_switch, visible_img_num, image_gallery, select_image, image_page_list],
1377
+ show_progress=opts.image_browser_show_progress
1378
+ ).then(
1379
+ fn=None,
1380
+ _js="image_browser_select_image",
1381
+ inputs=[tab_base_tag_box, image_index, select_image],
1382
+ outputs=[js_dummy_return],
1383
+ show_progress=opts.image_browser_show_progress
1384
+ )
1385
+
1386
+ to_dir_btn.click(save_image, inputs=[img_file_name, filenames, page_index, turn_page_switch, to_dir_path], outputs=[collected_warning, filenames, page_index, turn_page_switch], show_progress=opts.image_browser_show_progress)
1387
+ #turn page
1388
+ first_page.click(lambda s:(1, -s) , inputs=[turn_page_switch], outputs=[page_index, turn_page_switch], show_progress=opts.image_browser_show_progress)
1389
+ next_page.click(lambda p, s: (p + 1, -s), inputs=[page_index, turn_page_switch], outputs=[page_index, turn_page_switch], show_progress=opts.image_browser_show_progress)
1390
+ prev_page.click(lambda p, s: (p - 1, -s), inputs=[page_index, turn_page_switch], outputs=[page_index, turn_page_switch], show_progress=opts.image_browser_show_progress)
1391
+ end_page.click(lambda s: (-1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch], show_progress=opts.image_browser_show_progress)
1392
+ load_switch.change(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch], show_progress=opts.image_browser_show_progress)
1393
+ filename_keyword_search.submit(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch], show_progress=opts.image_browser_show_progress)
1394
+ exif_keyword_search.submit(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch], show_progress=opts.image_browser_show_progress)
1395
+ ranking_filter_min.submit(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch], show_progress=opts.image_browser_show_progress)
1396
+ ranking_filter_max.submit(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch], show_progress=opts.image_browser_show_progress)
1397
+ aes_filter_min.submit(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch], show_progress=opts.image_browser_show_progress)
1398
+ aes_filter_max.submit(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch], show_progress=opts.image_browser_show_progress)
1399
+ sort_by.change(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch], show_progress=opts.image_browser_show_progress)
1400
+ page_index.submit(lambda s: -s, inputs=[turn_page_switch], outputs=[turn_page_switch], show_progress=opts.image_browser_show_progress)
1401
+ renew_page.click(lambda s: -s, inputs=[turn_page_switch], outputs=[turn_page_switch], show_progress=opts.image_browser_show_progress)
1402
+ refresh_index_button.click(lambda p, s:(p, -s), inputs=[page_index, turn_page_switch], outputs=[page_index, turn_page_switch], show_progress=opts.image_browser_show_progress)
1403
+ img_path_depth.change(lambda s: -s, inputs=[turn_page_switch], outputs=[turn_page_switch], show_progress=opts.image_browser_show_progress)
1404
+
1405
+ hide_on_thumbnail_view = [delete_panel, button_panel, ranking_panel, to_dir_panel, info_add_panel]
1406
+
1407
+ sort_order.click(
1408
+ fn=sort_order_flip,
1409
+ inputs=[turn_page_switch, sort_order],
1410
+ outputs=[page_index, turn_page_switch, sort_order],
1411
+ show_progress=opts.image_browser_show_progress
1412
+ )
1413
+ ranking_filter.change(
1414
+ fn=ranking_filter_settings,
1415
+ inputs=[page_index, turn_page_switch, ranking_filter],
1416
+ outputs=[page_index, turn_page_switch, ranking_filter_min, ranking_filter_max],
1417
+ show_progress=opts.image_browser_show_progress
1418
+ )
1419
+
1420
+ # Others
1421
+ img_path_subdirs_button.click(
1422
+ fn=img_path_subdirs_get,
1423
+ inputs=[img_path],
1424
+ outputs=[img_path_subdirs],
1425
+ show_progress=opts.image_browser_show_progress
1426
+ )
1427
+ img_path_subdirs.change(
1428
+ fn=change_dir,
1429
+ inputs=[img_path_subdirs, path_recorder, load_switch, img_path_browser, img_path_depth, img_path],
1430
+ outputs=change_dir_outputs,
1431
+ show_progress=opts.image_browser_show_progress
1432
+ )
1433
+ img_path_save_button.click(
1434
+ fn=img_path_add_remove,
1435
+ inputs=[img_path, path_recorder, img_path_add, img_path_depth],
1436
+ outputs=[path_recorder, img_path_browser],
1437
+ show_progress=opts.image_browser_show_progress
1438
+ )
1439
+ img_path_remove_button.click(
1440
+ fn=img_path_add_remove,
1441
+ inputs=[img_path, path_recorder, img_path_remove, img_path_depth],
1442
+ outputs=[path_recorder, img_path_browser],
1443
+ show_progress=opts.image_browser_show_progress
1444
+ )
1445
+ maint_exif_rebuild.click(
1446
+ fn=exif_rebuild,
1447
+ inputs=[maint_wait],
1448
+ outputs=[maint_wait, maint_last_msg],
1449
+ show_progress=True
1450
+ )
1451
+ maint_exif_delete_0.click(
1452
+ fn=exif_delete_0,
1453
+ inputs=[maint_wait],
1454
+ outputs=[maint_wait, maint_last_msg],
1455
+ show_progress=True
1456
+ )
1457
+ maint_update_dirs.click(
1458
+ fn=exif_update_dirs,
1459
+ inputs=[maint_update_dirs_from, maint_update_dirs_to, maint_wait],
1460
+ outputs=[maint_wait, maint_last_msg],
1461
+ show_progress=True
1462
+ )
1463
+ maint_reapply_ranking.click(
1464
+ fn=reapply_ranking,
1465
+ inputs=[path_recorder, maint_wait],
1466
+ outputs=[maint_wait, maint_last_msg],
1467
+ show_progress=True
1468
+ )
1469
+ maint_restart_debug.click(
1470
+ fn=restart_debug,
1471
+ inputs=[maint_wait],
1472
+ outputs=[maint_wait, maint_last_msg],
1473
+ show_progress=True
1474
+ )
1475
+ maint_get_js_logs.click(
1476
+ fn=js_logs_output,
1477
+ _js="get_js_logs",
1478
+ inputs=[js_logs],
1479
+ outputs=[maint_show_logs],
1480
+ show_progress=True
1481
+ )
1482
+
1483
+ # other functions
1484
+ if opts.image_browser_use_thumbnail:
1485
+ set_index_outputs = [img_file_name, img_file_time, image_index, hidden, turn_page_switch, img_file_info_add, image_gallery]
1486
+ else:
1487
+ set_index_outputs = [img_file_name, img_file_time, image_index, hidden, turn_page_switch, img_file_info_add]
1488
+ set_index.click(
1489
+ fn=show_image_info,
1490
+ _js="image_browser_get_current_img",
1491
+ inputs=[tab_base_tag_box, image_index, page_index, filenames, turn_page_switch, image_gallery],
1492
+ outputs=set_index_outputs,
1493
+ show_progress=opts.image_browser_show_progress
1494
+ ).then(
1495
+ fn=None,
1496
+ _js="image_browser_img_show_progress_update",
1497
+ inputs=[],
1498
+ outputs=[js_dummy_return],
1499
+ show_progress=opts.image_browser_show_progress
1500
+ )
1501
+
1502
+ set_index.click(fn=lambda:(gr.update(visible=delete_panel not in override_hidden), gr.update(visible=button_panel not in override_hidden), gr.update(visible=ranking_panel not in override_hidden), gr.update(visible=to_dir_panel not in override_hidden), gr.update(visible=info_add_panel not in override_hidden)), inputs=None, outputs=hide_on_thumbnail_view, show_progress=opts.image_browser_show_progress)
1503
+
1504
+ favorites_btn.click(save_image, inputs=[img_file_name, filenames, page_index, turn_page_switch, favorites_path], outputs=[collected_warning, filenames, page_index, turn_page_switch], show_progress=opts.image_browser_show_progress)
1505
+ img_file_name.change(img_file_name_changed, inputs=[img_file_name, favorites_btn, to_dir_btn], outputs=[ranking_current, ranking, collected_warning, favorites_btn, to_dir_btn], show_progress=opts.image_browser_show_progress)
1506
+
1507
+ hidden.change(fn=run_pnginfo, inputs=[hidden, img_path, img_file_name], outputs=[info1, img_file_info, info2, image_browser_prompt, image_browser_neg_prompt], show_progress=opts.image_browser_show_progress)
1508
+
1509
+ #ranking
1510
+ ranking.change(update_ranking, inputs=[img_file_name, ranking_current, ranking, img_file_info], outputs=[ranking_current, ranking, img_file_info], show_progress=opts.image_browser_show_progress)
1511
+
1512
+ try:
1513
+ modules.generation_parameters_copypaste.bind_buttons(send_to_buttons, hidden, img_file_info)
1514
+ except:
1515
+ pass
1516
+
1517
+ if standard_ui:
1518
+ current_gr_tab.select(
1519
+ fn=tab_select,
1520
+ inputs=[],
1521
+ outputs=[path_recorder, to_dir_saved],
1522
+ show_progress=opts.image_browser_show_progress
1523
+ )
1524
+ open_folder_button.click(
1525
+ fn=lambda: open_folder(dir_name),
1526
+ inputs=[],
1527
+ outputs=[],
1528
+ show_progress=opts.image_browser_show_progress
1529
+ )
1530
+ elif others_dir:
1531
+ open_folder_button.click(
1532
+ fn=open_folder,
1533
+ inputs=[img_path],
1534
+ outputs=[],
1535
+ show_progress=opts.image_browser_show_progress
1536
+ )
1537
+ if standard_ui or others_dir:
1538
+ turn_page_switch.change(
1539
+ fn=get_image_page,
1540
+ inputs=[img_path, page_index, filenames, filename_keyword_search, sort_by, sort_order, tab_base_tag_box, img_path_depth, ranking_filter, ranking_filter_min, ranking_filter_max, aes_filter_min, aes_filter_max, score_type, exif_keyword_search, negative_prompt_search, use_regex, case_sensitive, image_reward_button],
1541
+ outputs=[filenames, page_index, image_gallery, img_file_name, img_file_time, img_file_info, visible_img_num, warning_box, hidden, image_page_list, image_browser_img_info, image_reward_button],
1542
+ show_progress=opts.image_browser_show_progress
1543
+ ).then(
1544
+ fn=None,
1545
+ _js="image_browser_turnpage",
1546
+ inputs=[tab_base_tag_box],
1547
+ outputs=[js_dummy_return],
1548
+ show_progress=opts.image_browser_show_progress
1549
+ )
1550
+ turn_page_switch.change(fn=lambda:(gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)), inputs=None, outputs=hide_on_thumbnail_view, show_progress=opts.image_browser_show_progress)
1551
+ sendto_openoutpaint.click(
1552
+ fn=None,
1553
+ inputs=[tab_base_tag_box, image_index, image_browser_prompt, image_browser_neg_prompt],
1554
+ outputs=[js_dummy_return],
1555
+ _js="image_browser_openoutpaint_send",
1556
+ show_progress=opts.image_browser_show_progress )
1557
+ sendto_controlnet_txt2img.click(
1558
+ fn=None,
1559
+ inputs=[tab_base_tag_box, image_index, sendto_controlnet_num, sendto_controlnet_type],
1560
+ outputs=[js_dummy_return],
1561
+ _js="image_browser_controlnet_send_txt2img",
1562
+ show_progress=opts.image_browser_show_progress
1563
+ )
1564
+ sendto_controlnet_img2img.click(
1565
+ fn=None,
1566
+ inputs=[tab_base_tag_box, image_index, sendto_controlnet_num, sendto_controlnet_type],
1567
+ outputs=[js_dummy_return],
1568
+ _js="image_browser_controlnet_send_img2img",
1569
+ show_progress=opts.image_browser_show_progress
1570
+ )
1571
+ image_reward_button.click(
1572
+ fn=generate_image_reward,
1573
+ inputs=[filenames, turn_page_switch, aes_filter_min, aes_filter_max],
1574
+ outputs=[turn_page_switch, aes_filter_min, aes_filter_max],
1575
+ show_progress=True
1576
+ )
1577
+
1578
+ def run_pnginfo(image, image_path, image_file_name):
1579
+ if image is None:
1580
+ return '', '', '', '', ''
1581
+ try:
1582
+ geninfo, items = images.read_info_from_image(image)
1583
+ items = {**{'parameters': geninfo}, **items}
1584
+
1585
+ info = ''
1586
+ for key, text in items.items():
1587
+ info += f"""
1588
+ <div>
1589
+ <p><b>{plaintext_to_html(str(key))}</b></p>
1590
+ <p>{plaintext_to_html(str(text))}</p>
1591
+ </div>
1592
+ """.strip()+"\n"
1593
+ except UnidentifiedImageError as e:
1594
+ geninfo = None
1595
+ info = ""
1596
+
1597
+ if geninfo is None:
1598
+ try:
1599
+ filename = os.path.splitext(image_file_name)[0] + ".txt"
1600
+ geninfo = ""
1601
+ with open(filename) as f:
1602
+ for line in f:
1603
+ geninfo += line
1604
+ except Exception:
1605
+ logger.warning(f"run_pnginfo: No EXIF in image or txt file")
1606
+
1607
+ if openoutpaint:
1608
+ prompt, neg_prompt = wib_db.select_prompts(image_file_name)
1609
+ if prompt == "0":
1610
+ prompt = ""
1611
+ if neg_prompt == "0":
1612
+ neg_prompt = ""
1613
+ else:
1614
+ prompt = ""
1615
+ neg_prompt = ""
1616
+
1617
+ return '', geninfo, info, prompt, neg_prompt
1618
+
1619
+
1620
+ def on_ui_tabs():
1621
+ global num_of_imgs_per_page, loads_files_num, js_dummy_return
1622
+ num_of_imgs_per_page = int(opts.image_browser_page_columns * opts.image_browser_page_rows)
1623
+ loads_files_num = int(opts.image_browser_pages_perload * num_of_imgs_per_page)
1624
+ with gr.Blocks(analytics_enabled=False) as image_browser:
1625
+ gradio_needed = "3.23.0"
1626
+ if version.parse(gr.__version__) < version.parse(gradio_needed):
1627
+ gr.HTML(f'<p style="color: red; font-weight: bold;">You are running Gradio version {gr.__version__}. This version of the extension requires at least Gradio version {gradio_needed}.</p><p style="color: red; font-weight: bold;">For more details see <a href="https://github.com/AlUlkesh/stable-diffusion-webui-images-browser/issues/116#issuecomment-1493259585" target="_blank">https://github.com/AlUlkesh/stable-diffusion-webui-images-browser/issues/116#issuecomment-1493259585</a></p>')
1628
+ else:
1629
+ with gr.Tabs(elem_id="image_browser_tabs_container") as tabs:
1630
+ js_dummy_return = gr.Textbox(interactive=False, visible=False)
1631
+ for i, tab in enumerate(tabs_list):
1632
+ with gr.Tab(tab.name, elem_id=f"{tab.base_tag}_image_browser_container") as current_gr_tab:
1633
+ with gr.Blocks(analytics_enabled=False):
1634
+ create_tab(tab, current_gr_tab)
1635
+ gr.Checkbox(value=opts.image_browser_preload, elem_id="image_browser_preload", visible=False)
1636
+ gr.Textbox(",".join( [tab.base_tag for tab in tabs_list] ), elem_id="image_browser_tab_base_tags_list", visible=False)
1637
+ gr.Checkbox(value=opts.image_browser_swipe, elem_id=f"image_browser_swipe")
1638
+
1639
+ javascript_level_value, (javascript_level, javascript_level_text) = debug_levels(arg_level="javascript")
1640
+ level_value, (level, level_text) = debug_levels(arg_text=opts.image_browser_debug_level)
1641
+ if level_value >= javascript_level_value:
1642
+ debug_level_option = level
1643
+ else:
1644
+ debug_level_option = ""
1645
+ gr.Textbox(value=debug_level_option, elem_id="image_browser_debug_level_option", visible=False)
1646
+
1647
+ return (image_browser, "Image Browser", "image_browser"),
1648
+
1649
+ def move_setting(cur_setting_name, old_setting_name, option_info, section, added):
1650
+ try:
1651
+ old_value = shared.opts.__getattr__(old_setting_name)
1652
+ except AttributeError:
1653
+ old_value = None
1654
+ try:
1655
+ new_value = shared.opts.__getattr__(cur_setting_name)
1656
+ except AttributeError:
1657
+ new_value = None
1658
+ if old_value is not None and new_value is None:
1659
+ # Add new option
1660
+ shared.opts.add_option(cur_setting_name, shared.OptionInfo(*option_info, section=section))
1661
+ shared.opts.__setattr__(cur_setting_name, old_value)
1662
+ added = added + 1
1663
+ # Remove old option
1664
+ shared.opts.data.pop(old_setting_name, None)
1665
+
1666
+ return added
1667
+
1668
+ def on_ui_settings():
1669
+ # [current setting_name], [old setting_name], [default], [label], [component], [component_args]
1670
+ active_tabs_description = f"List of active tabs (separated by commas). Available options are {', '.join(default_tab_options)}. Custom folders are also supported by specifying their path."
1671
+ debug_level_choices = []
1672
+ for i in range(len(debug_level_types)):
1673
+ level_value, (level, level_text) = debug_levels(arg_value=i)
1674
+ debug_level_choices.append(level_text)
1675
+
1676
+ image_browser_options = [
1677
+ ("image_browser_active_tabs", None, ", ".join(default_tab_options), active_tabs_description),
1678
+ ("image_browser_hidden_components", None, [], "Select components to hide", DropdownMulti, lambda: {"choices": components_list}),
1679
+ ("image_browser_with_subdirs", "images_history_with_subdirs", True, "Include images in sub directories"),
1680
+ ("image_browser_preload", "images_history_preload", False, "Preload images at startup for first tab"),
1681
+ ("image_browser_copy_image", "images_copy_image", False, "Move buttons copy instead of move"),
1682
+ ("image_browser_delete_message", "images_delete_message", True, "Print image deletion messages to the console"),
1683
+ ("image_browser_txt_files", "images_txt_files", True, "Move/Copy/Delete matching .txt files"),
1684
+ ("image_browser_debug_level", None, debug_level_choices[0], "Debug level", gr.Dropdown, lambda: {"choices": debug_level_choices}),
1685
+ ("image_browser_delete_recycle", "images_delete_recycle", True, "Use recycle bin when deleting images"),
1686
+ ("image_browser_scan_exif", "images_scan_exif", True, "Scan Exif-/.txt-data (initially slower, but required for many features to work)"),
1687
+ ("image_browser_mod_shift", None, False, "Change CTRL keybindings to SHIFT"),
1688
+ ("image_browser_mod_ctrl_shift", None, False, "or to CTRL+SHIFT"),
1689
+ ("image_browser_enable_maint", None, True, "Enable Maintenance tab"),
1690
+ ("image_browser_ranking_pnginfo", None, False, "Save ranking in image's pnginfo"),
1691
+ ("image_browser_page_columns", "images_history_page_columns", 6, "Number of columns on the page"),
1692
+ ("image_browser_page_rows", "images_history_page_rows", 6, "Number of rows on the page"),
1693
+ ("image_browser_pages_perload", "images_history_pages_perload", 20, "Minimum number of pages per load"),
1694
+ ("image_browser_use_thumbnail", None, False, "Use optimized images in the thumbnail interface (significantly reduces the amount of data transferred)"),
1695
+ ("image_browser_thumbnail_size", None, 200, "Size of the thumbnails (px)"),
1696
+ ("image_browser_swipe", None, False, "Swipe left/right navigates to the next image"),
1697
+ ("image_browser_img_tooltips", None, True, "Enable thumbnail tooltips"),
1698
+ ("image_browser_scoring_type", None, "aesthetic_score", "Default scoring type", gr.Dropdown, lambda: {"choices": ["aesthetic_score", "ImageReward Score"]}),
1699
+ ("image_browser_show_progress", None, True, "Show progress indicator"),
1700
+ ]
1701
+
1702
+ section = ('image-browser', "Image Browser")
1703
+ # Move historic setting names to current names
1704
+ added = 0
1705
+ for cur_setting_name, old_setting_name, *option_info in image_browser_options:
1706
+ if old_setting_name is not None:
1707
+ added = move_setting(cur_setting_name, old_setting_name, option_info, section, added)
1708
+ if added > 0:
1709
+ shared.opts.save(shared.config_filename)
1710
+
1711
+ for cur_setting_name, _, *option_info in image_browser_options:
1712
+ shared.opts.add_option(cur_setting_name, shared.OptionInfo(*option_info, section=section))
1713
+
1714
+ script_callbacks.on_ui_settings(on_ui_settings)
1715
+ script_callbacks.on_ui_tabs(on_ui_tabs)
scripts/wib/wib_db.py ADDED
@@ -0,0 +1,888 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import hashlib
2
+ import json
3
+ import os
4
+ import sqlite3
5
+ from modules import scripts
6
+ from PIL import Image
7
+
8
+ version = 6
9
+
10
+ path_recorder_file = os.path.join(scripts.basedir(), "path_recorder.txt")
11
+ aes_cache_file = os.path.join(scripts.basedir(), "aes_scores.json")
12
+ exif_cache_file = os.path.join(scripts.basedir(), "exif_data.json")
13
+ ranking_file = os.path.join(scripts.basedir(), "ranking.json")
14
+ archive = os.path.join(scripts.basedir(), "archive")
15
+ db_file = os.path.join(scripts.basedir(), "wib.sqlite3")
16
+ np = "Negative prompt: "
17
+ st = "Steps: "
18
+ timeout = 30
19
+
20
+ def create_filehash(cursor):
21
+ cursor.execute('''
22
+ CREATE TABLE IF NOT EXISTS filehash (
23
+ file TEXT PRIMARY KEY,
24
+ hash TEXT,
25
+ created TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
26
+ updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP
27
+ )
28
+ ''')
29
+
30
+ cursor.execute('''
31
+ CREATE TRIGGER filehash_tr
32
+ AFTER UPDATE ON filehash
33
+ BEGIN
34
+ UPDATE filehash SET updated = CURRENT_TIMESTAMP WHERE file = OLD.file;
35
+ END;
36
+ ''')
37
+
38
+ return
39
+
40
+ def create_work_files(cursor):
41
+ cursor.execute('''
42
+ CREATE TABLE IF NOT EXISTS work_files (
43
+ file TEXT PRIMARY KEY
44
+ )
45
+ ''')
46
+
47
+ return
48
+
49
+ def create_db(cursor):
50
+ cursor.execute('''
51
+ CREATE TABLE IF NOT EXISTS db_data (
52
+ key TEXT PRIMARY KEY,
53
+ value TEXT
54
+ )
55
+ ''')
56
+
57
+ cursor.execute('''
58
+ CREATE TABLE IF NOT EXISTS path_recorder (
59
+ path TEXT PRIMARY KEY,
60
+ depth INT,
61
+ path_display TEXT,
62
+ created TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
63
+ updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP
64
+ )
65
+ ''')
66
+
67
+ cursor.execute('''
68
+ CREATE TRIGGER path_recorder_tr
69
+ AFTER UPDATE ON path_recorder
70
+ BEGIN
71
+ UPDATE path_recorder SET updated = CURRENT_TIMESTAMP WHERE path = OLD.path;
72
+ END;
73
+ ''')
74
+
75
+ cursor.execute('''
76
+ CREATE TABLE IF NOT EXISTS exif_data (
77
+ file TEXT,
78
+ key TEXT,
79
+ value TEXT,
80
+ created TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
81
+ updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
82
+ PRIMARY KEY (file, key)
83
+ )
84
+ ''')
85
+
86
+ cursor.execute('''
87
+ CREATE INDEX IF NOT EXISTS exif_data_key ON exif_data (key)
88
+ ''')
89
+
90
+ cursor.execute('''
91
+ CREATE TRIGGER exif_data_tr
92
+ AFTER UPDATE ON exif_data
93
+ BEGIN
94
+ UPDATE exif_data SET updated = CURRENT_TIMESTAMP WHERE file = OLD.file AND key = OLD.key;
95
+ END;
96
+ ''')
97
+
98
+ cursor.execute('''
99
+ CREATE TABLE IF NOT EXISTS ranking (
100
+ file TEXT PRIMARY KEY,
101
+ name TEXT,
102
+ ranking TEXT,
103
+ created TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
104
+ updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP
105
+ )
106
+ ''')
107
+
108
+ cursor.execute('''
109
+ CREATE INDEX IF NOT EXISTS ranking_name ON ranking (name)
110
+ ''')
111
+
112
+ cursor.execute('''
113
+ CREATE TRIGGER ranking_tr
114
+ AFTER UPDATE ON ranking
115
+ BEGIN
116
+ UPDATE ranking SET updated = CURRENT_TIMESTAMP WHERE file = OLD.file;
117
+ END;
118
+ ''')
119
+
120
+ create_filehash(cursor)
121
+ create_work_files(cursor)
122
+
123
+ return
124
+
125
+ def migrate_path_recorder(cursor):
126
+ if os.path.exists(path_recorder_file):
127
+ try:
128
+ with open(path_recorder_file) as f:
129
+ # json-version
130
+ path_recorder = json.load(f)
131
+ for path, values in path_recorder.items():
132
+ path = os.path.realpath(path)
133
+ depth = values["depth"]
134
+ path_display = f"{path} [{depth}]"
135
+ cursor.execute('''
136
+ INSERT INTO path_recorder (path, depth, path_display)
137
+ VALUES (?, ?, ?)
138
+ ''', (path, depth, path_display))
139
+ except json.JSONDecodeError:
140
+ with open(path_recorder_file) as f:
141
+ # old txt-version
142
+ path = f.readline().rstrip("\n")
143
+ while len(path) > 0:
144
+ path = os.path.realpath(path)
145
+ cursor.execute('''
146
+ INSERT INTO path_recorder (path, depth, path_display)
147
+ VALUES (?, ?, ?)
148
+ ''', (path, 0, f"{path} [0]"))
149
+ path = f.readline().rstrip("\n")
150
+
151
+ return
152
+
153
+ def update_exif_data(cursor, file, info):
154
+ prompt = "0"
155
+ negative_prompt = "0"
156
+ key_values = "0: 0"
157
+ if info != "0":
158
+ info_list = info.split("\n")
159
+ prompt = ""
160
+ negative_prompt = ""
161
+ key_values = ""
162
+ for info_item in info_list:
163
+ if info_item.startswith(st):
164
+ key_values = info_item
165
+ elif info_item.startswith(np):
166
+ negative_prompt = info_item.replace(np, "")
167
+ else:
168
+ if prompt == "":
169
+ prompt = info_item
170
+ else:
171
+ # multiline prompts
172
+ prompt = f"{prompt}\n{info_item}"
173
+ if key_values != "":
174
+ key_value_pairs = []
175
+ key_value = ""
176
+ quote_open = False
177
+ for char in key_values + ",":
178
+ key_value += char
179
+ if char == '"':
180
+ quote_open = not quote_open
181
+ if char == "," and not quote_open:
182
+ try:
183
+ k, v = key_value.strip(" ,").split(": ")
184
+ except ValueError:
185
+ k = key_value.strip(" ,").split(": ")[0]
186
+ v = ""
187
+ key_value_pairs.append((k, v))
188
+ key_value = ""
189
+
190
+ try:
191
+ cursor.execute('''
192
+ INSERT INTO exif_data (file, key, value)
193
+ VALUES (?, ?, ?)
194
+ ''', (file, "prompt", prompt))
195
+ except sqlite3.IntegrityError:
196
+ # Duplicate, delete all "file" entries and try again
197
+ cursor.execute('''
198
+ DELETE FROM exif_data
199
+ WHERE file = ?
200
+ ''', (file,))
201
+
202
+ cursor.execute('''
203
+ INSERT INTO exif_data (file, key, value)
204
+ VALUES (?, ?, ?)
205
+ ''', (file, "prompt", prompt))
206
+
207
+ cursor.execute('''
208
+ INSERT INTO exif_data (file, key, value)
209
+ VALUES (?, ?, ?)
210
+ ''', (file, "negative_prompt", negative_prompt))
211
+
212
+ for (key, value) in key_value_pairs:
213
+ try:
214
+ cursor.execute('''
215
+ INSERT INTO exif_data (file, key, value)
216
+ VALUES (?, ?, ?)
217
+ ''', (file, key, value))
218
+ except sqlite3.IntegrityError:
219
+ pass
220
+
221
+ return
222
+
223
+ def migrate_exif_data(cursor):
224
+ if os.path.exists(exif_cache_file):
225
+ with open(exif_cache_file, 'r') as file:
226
+ exif_cache = json.load(file)
227
+
228
+ for file, info in exif_cache.items():
229
+ file = os.path.realpath(file)
230
+ update_exif_data(cursor, file, info)
231
+
232
+ return
233
+
234
+ def migrate_ranking(cursor):
235
+ if os.path.exists(ranking_file):
236
+ with open(ranking_file, 'r') as file:
237
+ ranking = json.load(file)
238
+ for file, info in ranking.items():
239
+ if info != "None":
240
+ file = os.path.realpath(file)
241
+ name = os.path.basename(file)
242
+ cursor.execute('''
243
+ INSERT INTO ranking (file, name, ranking)
244
+ VALUES (?, ?, ?)
245
+ ''', (file, name, info))
246
+
247
+ return
248
+
249
+ def get_hash(file):
250
+ # Get filehash without exif info
251
+ try:
252
+ image = Image.open(file)
253
+ except Exception as e:
254
+ print(e)
255
+
256
+ hash = hashlib.sha512(image.tobytes()).hexdigest()
257
+ image.close()
258
+
259
+ return hash
260
+
261
+ def migrate_filehash(cursor, version):
262
+ if version <= "4":
263
+ create_filehash(cursor)
264
+
265
+ cursor.execute('''
266
+ SELECT file
267
+ FROM ranking
268
+ ''')
269
+ for (file,) in cursor.fetchall():
270
+ if os.path.exists(file):
271
+ hash = get_hash(file)
272
+ cursor.execute('''
273
+ INSERT INTO filehash (file, hash)
274
+ VALUES (?, ?)
275
+ ''', (file, hash))
276
+
277
+ return
278
+
279
+ def migrate_work_files(cursor):
280
+ create_work_files(cursor)
281
+
282
+ return
283
+
284
+ def update_db_data(cursor, key, value):
285
+ cursor.execute('''
286
+ INSERT OR REPLACE
287
+ INTO db_data (key, value)
288
+ VALUES (?, ?)
289
+ ''', (key, value))
290
+
291
+ return
292
+
293
+ def get_version():
294
+ with sqlite3.connect(db_file, timeout=timeout) as conn:
295
+ cursor = conn.cursor()
296
+ cursor.execute('''
297
+ SELECT value
298
+ FROM db_data
299
+ WHERE key = 'version'
300
+ ''',)
301
+ db_version = cursor.fetchone()
302
+
303
+ return db_version
304
+
305
+ def migrate_path_recorder_dirs(cursor):
306
+ cursor.execute('''
307
+ SELECT path, path_display
308
+ FROM path_recorder
309
+ ''')
310
+ for (path, path_display) in cursor.fetchall():
311
+ real_path = os.path.realpath(path)
312
+ if path != real_path:
313
+ update_from = path
314
+ update_to = real_path
315
+ try:
316
+ cursor.execute('''
317
+ UPDATE path_recorder
318
+ SET path = ?,
319
+ path_display = ? || SUBSTR(path_display, LENGTH(?) + 1)
320
+ WHERE path = ?
321
+ ''', (update_to, update_to, update_from, update_from))
322
+ except sqlite3.IntegrityError as e:
323
+ # these are double keys, because the same file can be in the db with different path notations
324
+ (e_msg,) = e.args
325
+ if e_msg.startswith("UNIQUE constraint"):
326
+ cursor.execute('''
327
+ DELETE FROM path_recorder
328
+ WHERE path = ?
329
+ ''', (update_from,))
330
+ else:
331
+ raise
332
+
333
+ return
334
+
335
+ def migrate_exif_data_dirs(cursor):
336
+ cursor.execute('''
337
+ SELECT file
338
+ FROM exif_data
339
+ ''')
340
+ for (filepath,) in cursor.fetchall():
341
+ (path, file) = os.path.split(filepath)
342
+ real_path = os.path.realpath(path)
343
+ if path != real_path:
344
+ update_from = filepath
345
+ update_to = os.path.join(real_path, file)
346
+ try:
347
+ cursor.execute('''
348
+ UPDATE exif_data
349
+ SET file = ?
350
+ WHERE file = ?
351
+ ''', (update_to, update_from))
352
+ except sqlite3.IntegrityError as e:
353
+ # these are double keys, because the same file can be in the db with different path notations
354
+ (e_msg,) = e.args
355
+ if e_msg.startswith("UNIQUE constraint"):
356
+ cursor.execute('''
357
+ DELETE FROM exif_data
358
+ WHERE file = ?
359
+ ''', (update_from,))
360
+ else:
361
+ raise
362
+
363
+ return
364
+
365
+ def migrate_ranking_dirs(cursor, db_version):
366
+ if db_version == "1":
367
+ cursor.execute('''
368
+ ALTER TABLE ranking
369
+ ADD COLUMN name TEXT
370
+ ''')
371
+
372
+ cursor.execute('''
373
+ CREATE INDEX IF NOT EXISTS ranking_name ON ranking (name)
374
+ ''')
375
+
376
+ cursor.execute('''
377
+ SELECT file, ranking
378
+ FROM ranking
379
+ ''')
380
+ for (filepath, ranking) in cursor.fetchall():
381
+ if filepath == "" or ranking == "None":
382
+ cursor.execute('''
383
+ DELETE FROM ranking
384
+ WHERE file = ?
385
+ ''', (filepath,))
386
+ else:
387
+ (path, file) = os.path.split(filepath)
388
+ real_path = os.path.realpath(path)
389
+ name = file
390
+ update_from = filepath
391
+ update_to = os.path.join(real_path, file)
392
+ try:
393
+ cursor.execute('''
394
+ UPDATE ranking
395
+ SET file = ?,
396
+ name = ?
397
+ WHERE file = ?
398
+ ''', (update_to, name, update_from))
399
+ except sqlite3.IntegrityError as e:
400
+ # these are double keys, because the same file can be in the db with different path notations
401
+ (e_msg,) = e.args
402
+ if e_msg.startswith("UNIQUE constraint"):
403
+ cursor.execute('''
404
+ DELETE FROM ranking
405
+ WHERE file = ?
406
+ ''', (update_from,))
407
+ else:
408
+ raise
409
+
410
+ return
411
+
412
+ def check():
413
+ if not os.path.exists(db_file):
414
+ conn, cursor = transaction_begin()
415
+ print("Image Browser: Creating database")
416
+ create_db(cursor)
417
+ update_db_data(cursor, "version", version)
418
+ migrate_path_recorder(cursor)
419
+ migrate_exif_data(cursor)
420
+ migrate_ranking(cursor)
421
+ migrate_filehash(cursor, str(version))
422
+ transaction_end(conn, cursor)
423
+ print("Image Browser: Database created")
424
+ db_version = get_version()
425
+ conn, cursor = transaction_begin()
426
+ if db_version[0] <= "2":
427
+ # version 1 database had mixed path notations, changed them all to abspath
428
+ # version 2 database still had mixed path notations, because of windows short name, changed them all to realpath
429
+ print(f"Image Browser: Upgrading database from version {db_version[0]} to version {version}")
430
+ migrate_path_recorder_dirs(cursor)
431
+ migrate_exif_data_dirs(cursor)
432
+ migrate_ranking_dirs(cursor, db_version[0])
433
+ if db_version[0] <= "4":
434
+ migrate_filehash(cursor, db_version[0])
435
+ if db_version[0] <= "5":
436
+ migrate_work_files(cursor)
437
+ update_db_data(cursor, "version", version)
438
+ print(f"Image Browser: Database upgraded from version {db_version[0]} to version {version}")
439
+ transaction_end(conn, cursor)
440
+
441
+ return version
442
+
443
+ def load_path_recorder():
444
+ with sqlite3.connect(db_file, timeout=timeout) as conn:
445
+ cursor = conn.cursor()
446
+ cursor.execute('''
447
+ SELECT path, depth, path_display
448
+ FROM path_recorder
449
+ ''')
450
+ path_recorder = {path: {"depth": depth, "path_display": path_display} for path, depth, path_display in cursor.fetchall()}
451
+
452
+ return path_recorder
453
+
454
+ def select_ranking(file):
455
+ with sqlite3.connect(db_file, timeout=timeout) as conn:
456
+ cursor = conn.cursor()
457
+ cursor.execute('''
458
+ SELECT ranking
459
+ FROM ranking
460
+ WHERE file = ?
461
+ ''', (file,))
462
+ ranking_value = cursor.fetchone()
463
+
464
+ if ranking_value is None:
465
+ return_ranking = "None"
466
+ else:
467
+ (return_ranking,) = ranking_value
468
+
469
+ return return_ranking
470
+
471
+ def update_ranking(file, ranking):
472
+ name = os.path.basename(file)
473
+ with sqlite3.connect(db_file, timeout=timeout) as conn:
474
+ cursor = conn.cursor()
475
+ if ranking == "None":
476
+ cursor.execute('''
477
+ DELETE FROM ranking
478
+ WHERE file = ?
479
+ ''', (file,))
480
+ else:
481
+ cursor.execute('''
482
+ INSERT OR REPLACE
483
+ INTO ranking (file, name, ranking)
484
+ VALUES (?, ?, ?)
485
+ ''', (file, name, ranking))
486
+
487
+ hash = get_hash(file)
488
+ cursor.execute('''
489
+ INSERT OR REPLACE
490
+ INTO filehash (file, hash)
491
+ VALUES (?, ?)
492
+ ''', (file, hash))
493
+
494
+ return
495
+
496
+ def select_image_reward_score(cursor, file):
497
+ cursor.execute('''
498
+ SELECT value
499
+ FROM exif_data
500
+ WHERE file = ?
501
+ AND key = 'ImageRewardScore'
502
+ ''', (file,))
503
+ image_reward_score = cursor.fetchone()
504
+ if image_reward_score is None:
505
+ return_image_reward_score = None
506
+ else:
507
+ (return_image_reward_score,) = image_reward_score
508
+ cursor.execute('''
509
+ SELECT value
510
+ FROM exif_data
511
+ WHERE file = ?
512
+ AND key = 'prompt'
513
+ ''', (file,))
514
+ image_reward_prompt = cursor.fetchone()
515
+ if image_reward_prompt is None:
516
+ return_image_reward_prompt = None
517
+ else:
518
+ (return_image_reward_prompt,) = image_reward_prompt
519
+
520
+ return return_image_reward_score, return_image_reward_prompt
521
+
522
+ def update_image_reward_score(cursor, file, image_reward_score):
523
+ cursor.execute('''
524
+ INSERT OR REPLACE
525
+ INTO exif_data (file, key, value)
526
+ VALUES (?, ?, ?)
527
+ ''', (file, "ImageRewardScore", image_reward_score))
528
+
529
+ return
530
+
531
+ def update_path_recorder(path, depth, path_display):
532
+ with sqlite3.connect(db_file, timeout=timeout) as conn:
533
+ cursor = conn.cursor()
534
+ cursor.execute('''
535
+ INSERT OR REPLACE
536
+ INTO path_recorder (path, depth, path_display)
537
+ VALUES (?, ?, ?)
538
+ ''', (path, depth, path_display))
539
+
540
+ return
541
+
542
+ def update_path_recorder(path, depth, path_display):
543
+ with sqlite3.connect(db_file, timeout=timeout) as conn:
544
+ cursor = conn.cursor()
545
+ cursor.execute('''
546
+ INSERT OR REPLACE
547
+ INTO path_recorder (path, depth, path_display)
548
+ VALUES (?, ?, ?)
549
+ ''', (path, depth, path_display))
550
+
551
+ return
552
+
553
+ def delete_path_recorder(path):
554
+ with sqlite3.connect(db_file, timeout=timeout) as conn:
555
+ cursor = conn.cursor()
556
+ cursor.execute('''
557
+ DELETE FROM path_recorder
558
+ WHERE path = ?
559
+ ''', (path,))
560
+
561
+ return
562
+
563
+ def update_path_recorder_mult(cursor, update_from, update_to):
564
+ cursor.execute('''
565
+ UPDATE path_recorder
566
+ SET path = ?,
567
+ path_display = ? || SUBSTR(path_display, LENGTH(?) + 1)
568
+ WHERE path = ?
569
+ ''', (update_to, update_to, update_from, update_from))
570
+
571
+ return
572
+
573
+ def update_exif_data_mult(cursor, update_from, update_to):
574
+ update_from = update_from + os.path.sep
575
+ update_to = update_to + os.path.sep
576
+ cursor.execute('''
577
+ UPDATE exif_data
578
+ SET file = ? || SUBSTR(file, LENGTH(?) + 1)
579
+ WHERE file like ? || '%'
580
+ ''', (update_to, update_from, update_from))
581
+
582
+ return
583
+
584
+ def update_ranking_mult(cursor, update_from, update_to):
585
+ update_from = update_from + os.path.sep
586
+ update_to = update_to + os.path.sep
587
+ cursor.execute('''
588
+ UPDATE ranking
589
+ SET file = ? || SUBSTR(file, LENGTH(?) + 1)
590
+ WHERE file like ? || '%'
591
+ ''', (update_to, update_from, update_from))
592
+
593
+ return
594
+
595
+ def delete_exif_0(cursor):
596
+ cursor.execute('''
597
+ DELETE FROM exif_data
598
+ WHERE file IN (
599
+ SELECT file FROM exif_data a
600
+ WHERE value = '0'
601
+ GROUP BY file
602
+ HAVING COUNT(*) = (SELECT COUNT(*) FROM exif_data WHERE file = a.file)
603
+ )
604
+ ''')
605
+
606
+ return
607
+
608
+ def get_ranking_by_file(cursor, file):
609
+ cursor.execute('''
610
+ SELECT ranking
611
+ FROM ranking
612
+ WHERE file = ?
613
+ ''', (file,))
614
+ ranking_value = cursor.fetchone()
615
+
616
+ return ranking_value
617
+
618
+ def get_ranking_by_name(cursor, name):
619
+ cursor.execute('''
620
+ SELECT file, ranking
621
+ FROM ranking
622
+ WHERE name = ?
623
+ ''', (name,))
624
+ ranking_value = cursor.fetchone()
625
+
626
+ if ranking_value is not None:
627
+ (file, _) = ranking_value
628
+ cursor.execute('''
629
+ SELECT hash
630
+ FROM filehash
631
+ WHERE file = ?
632
+ ''', (file,))
633
+ hash_value = cursor.fetchone()
634
+ else:
635
+ hash_value = None
636
+
637
+ return ranking_value, hash_value
638
+
639
+ def insert_ranking(cursor, file, ranking, hash):
640
+ name = os.path.basename(file)
641
+ cursor.execute('''
642
+ INSERT INTO ranking (file, name, ranking)
643
+ VALUES (?, ?, ?)
644
+ ''', (file, name, ranking))
645
+
646
+ cursor.execute('''
647
+ INSERT OR REPLACE
648
+ INTO filehash (file, hash)
649
+ VALUES (?, ?)
650
+ ''', (file, hash))
651
+
652
+ return
653
+
654
+ def replace_ranking(cursor, file, alternate_file, hash):
655
+ name = os.path.basename(file)
656
+ cursor.execute('''
657
+ UPDATE ranking
658
+ SET file = ?
659
+ WHERE file = ?
660
+ ''', (file, alternate_file))
661
+
662
+ cursor.execute('''
663
+ INSERT OR REPLACE
664
+ INTO filehash (file, hash)
665
+ VALUES (?, ?)
666
+ ''', (file, hash))
667
+
668
+ return
669
+
670
+ def transaction_begin():
671
+ conn = sqlite3.connect(db_file, timeout=timeout)
672
+ conn.isolation_level = None
673
+ cursor = conn.cursor()
674
+ cursor.execute("BEGIN")
675
+ return conn, cursor
676
+
677
+ def transaction_end(conn, cursor):
678
+ cursor.execute("COMMIT")
679
+ conn.close()
680
+ return
681
+
682
+ def update_exif_data_by_key(cursor, file, key, value):
683
+ cursor.execute('''
684
+ INSERT OR REPLACE
685
+ INTO exif_data (file, key, value)
686
+ VALUES (?, ?, ?)
687
+ ''', (file, key, value))
688
+
689
+ return
690
+
691
+ def select_prompts(file):
692
+ with sqlite3.connect(db_file, timeout=timeout) as conn:
693
+ cursor = conn.cursor()
694
+ cursor.execute('''
695
+ SELECT key, value
696
+ FROM exif_data
697
+ WHERE file = ?
698
+ AND KEY in ('prompt', 'negative_prompt')
699
+ ''', (file,))
700
+
701
+ rows = cursor.fetchall()
702
+ prompt = ""
703
+ neg_prompt = ""
704
+ for row in rows:
705
+ (key, value) = row
706
+ if key == 'prompt':
707
+ prompt = value
708
+ elif key == 'negative_prompt':
709
+ neg_prompt = value
710
+
711
+ return prompt, neg_prompt
712
+
713
+ def load_exif_data(exif_cache):
714
+ with sqlite3.connect(db_file, timeout=timeout) as conn:
715
+ cursor = conn.cursor()
716
+ cursor.execute('''
717
+ SELECT file, group_concat(
718
+ case when key = 'prompt' or key = 'negative_prompt' then key || ': ' || value || '\n'
719
+ else key || ': ' || value
720
+ end, ', ') AS string
721
+ FROM (
722
+ SELECT *
723
+ FROM exif_data
724
+ ORDER BY
725
+ CASE WHEN key = 'prompt' THEN 0
726
+ WHEN key = 'negative_prompt' THEN 1
727
+ ELSE 2 END,
728
+ key
729
+ )
730
+ GROUP BY file
731
+ ''')
732
+
733
+ rows = cursor.fetchall()
734
+ for row in rows:
735
+ exif_cache[row[0]] = row[1]
736
+
737
+ return exif_cache
738
+
739
+ def load_exif_data_by_key(cache, key1, key2):
740
+ with sqlite3.connect(db_file, timeout=timeout) as conn:
741
+ cursor = conn.cursor()
742
+ cursor.execute('''
743
+ SELECT file, value
744
+ FROM exif_data
745
+ WHERE key IN (?, ?)
746
+ ''', (key1, key2))
747
+
748
+ rows = cursor.fetchall()
749
+ for row in rows:
750
+ cache[row[0]] = row[1]
751
+
752
+ return cache
753
+
754
+ def get_exif_dirs():
755
+ with sqlite3.connect(db_file, timeout=timeout) as conn:
756
+ cursor = conn.cursor()
757
+ cursor.execute('''
758
+ SELECT file
759
+ FROM exif_data
760
+ ''')
761
+
762
+ rows = cursor.fetchall()
763
+
764
+ dirs = {}
765
+ for row in rows:
766
+ dir = os.path.dirname(row[0])
767
+ dirs[dir] = dir
768
+
769
+ return dirs
770
+
771
+ def fill_work_files(cursor, fileinfos):
772
+ filenames = [x[0] for x in fileinfos]
773
+
774
+ cursor.execute('''
775
+ DELETE
776
+ FROM work_files
777
+ ''')
778
+
779
+ sql = '''
780
+ INSERT INTO work_files (file)
781
+ VALUES (?)
782
+ '''
783
+
784
+ cursor.executemany(sql, [(x,) for x in filenames])
785
+
786
+ return
787
+
788
+ def filter_aes(cursor, fileinfos, aes_filter_min_num, aes_filter_max_num, score_type):
789
+ if score_type == "aesthetic_score":
790
+ key = "aesthetic_score"
791
+ else:
792
+ key = "ImageRewardScore"
793
+
794
+ cursor.execute('''
795
+ DELETE
796
+ FROM work_files
797
+ WHERE file not in (
798
+ SELECT file
799
+ FROM exif_data b
800
+ WHERE file = b.file
801
+ AND b.key = ?
802
+ AND CAST(b.value AS REAL) between ? and ?
803
+ )
804
+ ''', (key, aes_filter_min_num, aes_filter_max_num))
805
+
806
+ cursor.execute('''
807
+ SELECT file
808
+ FROM work_files
809
+ ''')
810
+
811
+ rows = cursor.fetchall()
812
+
813
+ fileinfos_dict = {pair[0]: pair[1] for pair in fileinfos}
814
+ fileinfos_new = []
815
+ for (file,) in rows:
816
+ if fileinfos_dict.get(file) is not None:
817
+ fileinfos_new.append((file, fileinfos_dict[file]))
818
+
819
+ return fileinfos_new
820
+
821
+ def filter_ranking(cursor, fileinfos, ranking_filter, ranking_filter_min_num, ranking_filter_max_num):
822
+ if ranking_filter == "None":
823
+ cursor.execute('''
824
+ DELETE
825
+ FROM work_files
826
+ WHERE file IN (
827
+ SELECT file
828
+ FROM ranking b
829
+ WHERE file = b.file
830
+ )
831
+ ''')
832
+ elif ranking_filter == "Min-max":
833
+ cursor.execute('''
834
+ DELETE
835
+ FROM work_files
836
+ WHERE file NOT IN (
837
+ SELECT file
838
+ FROM ranking b
839
+ WHERE file = b.file
840
+ AND b.ranking BETWEEN ? AND ?
841
+ )
842
+ ''', (ranking_filter_min_num, ranking_filter_max_num))
843
+ else:
844
+ cursor.execute('''
845
+ DELETE
846
+ FROM work_files
847
+ WHERE file NOT IN (
848
+ SELECT file
849
+ FROM ranking b
850
+ WHERE file = b.file
851
+ AND b.ranking = ?
852
+ )
853
+ ''', (ranking_filter,))
854
+
855
+ cursor.execute('''
856
+ SELECT file
857
+ FROM work_files
858
+ ''')
859
+
860
+ rows = cursor.fetchall()
861
+
862
+ fileinfos_dict = {pair[0]: pair[1] for pair in fileinfos}
863
+ fileinfos_new = []
864
+ for (file,) in rows:
865
+ if fileinfos_dict.get(file) is not None:
866
+ fileinfos_new.append((file, fileinfos_dict[file]))
867
+
868
+ return fileinfos_new
869
+
870
+ def select_x_y(cursor, file):
871
+ cursor.execute('''
872
+ SELECT value
873
+ FROM exif_data
874
+ WHERE file = ?
875
+ AND key = 'Size'
876
+ ''', (file,))
877
+ size_value = cursor.fetchone()
878
+
879
+ if size_value is None:
880
+ x = "?"
881
+ y = "?"
882
+ else:
883
+ (size,) = size_value
884
+ parts = size.split("x")
885
+ x = parts[0]
886
+ y = parts[1]
887
+
888
+ return x, y
style.css ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ .thumbnails.svelte-1tkea93.svelte-1tkea93 {
2
+ justify-content: initial;
3
+ }
4
+
5
+ .thumbnails.scroll-hide.svelte-g4rw9 {
6
+ justify-content: initial;
7
+ }
8
+
9
+ div[id^="image_browser_tab"][id$="image_browser_gallery"].hide_loading > .svelte-gjihhp {
10
+ display: none;
11
+ }
12
+
13
+ .image_browser_gallery img {
14
+ object-fit: scale-down !important;
15
+ }