try streaming dataset

app.py

@@ -55,7 +55,8 @@ class GalleryApp:
                     if idx + j < len(items):
                         with cols[j]:
                             # show image
-                            image = st.session_state.images[items.iloc[idx+j]['row_idx'].item()]['image']
+                            # image = st.session_state.images[items.iloc[idx+j]['row_idx'].item()]['image']
+                            image = list(st.session_state.images.skip(items.iloc[idx+j]['row_idx'].item()).take(1))[0]['image']
                             st.image(image,
                                      use_column_width=True,
                             )
@@ -210,7 +211,8 @@ if __name__ == '__main__':
         if 'checked' not in st.session_state.promptBook.columns:
             st.session_state.promptBook['checked'] = False
 
-        st.session_state.images = load_from_disk(os.path.join(os.getcwd(), 'data', 'promptbook'))
+        # st.session_state.images = load_from_disk(os.path.join(os.getcwd(), 'data', 'promptbook'))
+        st.session_state.images = load_dataset('NYUSHPRP/ModelCofferPromptBook', split='train', streaming=True)
         print(st.session_state.images)
         print('images loaded')
         # st.session_state.promptBook = pd.DataFrame(load_dataset('NYUSHPRP/ModelCofferPromptBook', split='train'))