no more globals, iterator is by user session. It was generating some conflicts

app.py

@@ -23,10 +23,7 @@ sets = {
     }
 }
 
-def open_dataset(dataset, set_name, split, batch_size, shard = -1):
-    # I should really move ds.config_name and dsi to a gr.State()
-    global dsi, ds
-
+def open_dataset(dataset, set_name, split, batch_size, state, shard = -1):
     if shard == -1:
         # Trick to open the whole dataset
         data_files = None
@@ -53,14 +50,15 @@ def open_dataset(dataset, set_name, split, batch_size, shard = -1):
     
         dsi = iter(ds)
 
+    state["config"]  = config
+    state["dsi"] = dsi
     return (
         gr.update(label=f"Shards (max {shards})", value=shard, maximum=shards),
-        *get_images(batch_size)
+        *get_images(batch_size, state),
+        state
     )
     
-def get_images(batch_size):
-    global dsi, ds
-        
+def get_images(batch_size, state):
     items = []
     metadatas = []
 
@@ -70,17 +68,17 @@ def get_images(batch_size):
             metadata = {"bounds":[[1,1,4,4]], }
         else:
             try:
-                item = next(dsi)
+                item = next(state["dsi"])
             except StopIteration:
                 break
             metadata = item["metadata"]
-            if ds.config_name == "satellogic":
+            if state["config"] == "satellogic":
                 # image = (np.asarray(item["1m"])).astype("uint8")
                 # items.append(image[0,0,:,:])
                 image = np.asarray(item["rgb"][0]).astype(np.uint8)
                 items.append(image.transpose(1,2,0))
 
-            if ds.config_name == "sentinel_1":
+            if state["config"] == "sentinel_1":
                 metadata = json.loads(metadata)
                 data = np.asarray(item["10m"])
                 for i in range(data.shape[0]):
@@ -92,7 +90,7 @@ def get_images(batch_size):
                     image[2] = (image[0]/(image[1]+0.1))*256
                     items.append(image.transpose(1,2,0))
 
-            if ds.config_name == "default":
+            if state["config"] == "default":
                 dataRGB = np.asarray(item["rgb"]).astype("uint8")
                 dataCHM = np.asarray(item["chm"]).astype("uint8")
                 data1m  = np.asarray(item["1m"]).astype("uint8")
@@ -112,9 +110,13 @@ def get_images(batch_size):
 def update_shape(rows, columns):
     return gr.update(rows=rows, columns=columns)
 
+def new_state():
+    return gr.State({})
 
 with gr.Blocks(title="Dataset Explorer", fill_height = True) as demo:
-    gr.Markdown("# [satellogic/EarthView](https://huggingface.co/datasets/satellogic/EarthView) Dataset Viewer")
+    state = new_state()
+
+    gr.Markdown("# Viewer for [satellogic/EarthView](https://huggingface.co/datasets/satellogic/EarthView) Dataset")
     batch_size = gr.Number(10, label = "Batch Size", render=False)
     shard = gr.Slider(label="Shard", minimum=0, maximum=10000, step=1, render=False)
     table = gr.DataFrame(render = False)
@@ -133,8 +135,8 @@ with gr.Blocks(title="Dataset Explorer", fill_height = True) as demo:
 
         gr.Button("Load (minutes)").click(
             open_dataset,
-            inputs=[dataset, config, split, batch_size, initial_shard],
-            outputs=[shard, gallery, table])
+            inputs=[dataset, config, split, batch_size, state, initial_shard],
+            outputs=[shard, gallery, table, state])
 
     gallery.render()
     
@@ -151,11 +153,11 @@ with gr.Blocks(title="Dataset Explorer", fill_height = True) as demo:
         shard.render()
         shard.release(
             open_dataset,
-            inputs=[dataset, config, split, batch_size, shard],
-            outputs=[shard, gallery, table])
+            inputs=[dataset, config, split, batch_size, state, shard],
+            outputs=[shard, gallery, table, state])
 
-        btn = gr.Button("Get More Images", scale=0)
-        btn.click(get_images, [batch_size], [gallery, table])
+        btn = gr.Button("Next Batch (same shard)", scale=0)
+        btn.click(get_images, [batch_size, state], [gallery, table])
         btn.click()
     
     table.render()