1 files changed, 31 insertions, 5 deletions
diff --git a/javascript/ui.js b/javascript/ui.js
index 7aa30dc1..a73eeaa2 100644
--- a/javascript/ui.js
+++ b/javascript/ui.js
@@ -7,9 +7,31 @@ function set_theme(theme){
     }
 }
 
+function all_gallery_buttons() {
+    var allGalleryButtons = gradioApp().querySelectorAll('[style="display: block;"].tabitem div[id$=_gallery].gradio-gallery .thumbnails > .thumbnail-item.thumbnail-small');
+    var visibleGalleryButtons = [];
+    allGalleryButtons.forEach(function(elem) {
+        if (elem.parentElement.offsetParent) {
+            visibleGalleryButtons.push(elem);
+        }
+    })
+    return visibleGalleryButtons;
+}
+
+function selected_gallery_button() {
+    var allCurrentButtons = gradioApp().querySelectorAll('[style="display: block;"].tabitem div[id$=_gallery].gradio-gallery .thumbnail-item.thumbnail-small.selected');
+    var visibleCurrentButton = null;
+    allCurrentButtons.forEach(function(elem) {
+        if (elem.parentElement.offsetParent) {
+            visibleCurrentButton = elem;
+        }
+    })
+    return visibleCurrentButton;
+}
+
 function selected_gallery_index(){
-    var buttons = gradioApp().querySelectorAll('[style="display: block;"].tabitem div[id$=_gallery] .thumbnails > .thumbnail-item')
-    var button = gradioApp().querySelector('[style="display: block;"].tabitem div[id$=_gallery] .thumbnails > .thumbnail-item.selected')
+    var buttons = all_gallery_buttons();
+    var button = selected_gallery_button();
 
     var result = -1
     buttons.forEach(function(v, i){ if(v==button) { result = i } })
@@ -18,14 +40,18 @@ function selected_gallery_index(){
 }
 
 function extract_image_from_gallery(gallery){
-    if(gallery.length == 1){
-        return [gallery[0]]
+    if (gallery.length == 0){
+        return [null];
+    }
+    if (gallery.length == 1){
+        return [gallery[0]];
     }
 
     index = selected_gallery_index()
 
     if (index < 0 || index >= gallery.length){
-        return [null]
+        // Use the first image in the gallery as the default
+        index = 0;
     }
 
     return [gallery[index]];