Update app.py

app.py

@@ -100,6 +100,9 @@ with gr.Blocks() as demo:
     dataset_html = gr.HTML()
     message_box = gr.Textbox(interactive=False, label="Message")
 
+    # Entry selector initialized here
+    entry_selector = gr.Dropdown(label="Select Entry to Edit/Delete")
+
     with gr.Tab("Create / Upload Dataset"):
         with gr.Row():
             with gr.Column():
@@ -153,7 +156,7 @@ with gr.Blocks() as demo:
     dataset_selector.change(
         select_dataset,
         inputs=[dataset_selector, datasets],
-        outputs=[current_dataset_name, dataset_html, gr.Variable(), message_box]
+        outputs=[current_dataset_name, dataset_html, entry_selector, message_box]
     )
 
     with gr.Tab("Add Entry"):
@@ -183,18 +186,20 @@ with gr.Blocks() as demo:
     add_button.click(
         add_entry,
         inputs=[image_input, prompt_input, current_dataset_name, datasets],
-        outputs=[datasets, dataset_html, gr.Variable(), message_box]
+        outputs=[datasets, dataset_html, entry_selector, message_box]
     )
 
     with gr.Tab("Edit / Delete Entry"):
-        # Entry selector
-        entry_selector = gr.Dropdown(label="Select Entry to Edit/Delete")
-        selected_image = gr.Image(label="Selected Image", interactive=False)
-        selected_prompt = gr.Textbox(label="Current Prompt", interactive=False)
-        new_prompt_input = gr.Textbox(label="New Prompt (for Edit)")
-        with gr.Row():
-            edit_button = gr.Button("Edit Entry")
-            delete_button = gr.Button("Delete Entry")
+        # Components moved inside the tab
+        with gr.Column():
+            # Entry selector
+            entry_selector = gr.Dropdown(label="Select Entry to Edit/Delete")
+            selected_image = gr.Image(label="Selected Image", interactive=False)
+            selected_prompt = gr.Textbox(label="Current Prompt", interactive=False)
+            new_prompt_input = gr.Textbox(label="New Prompt (for Edit)")
+            with gr.Row():
+                edit_button = gr.Button("Edit Entry")
+                delete_button = gr.Button("Delete Entry")
 
     def update_selected_entry(entry_option, current_dataset_name, datasets):
         if not current_dataset_name or not entry_option:
@@ -203,7 +208,11 @@ with gr.Blocks() as demo:
         entry = datasets[current_dataset_name][index]
         image_data = entry['image']
         prompt = entry['prompt']
-        return gr.update(value=image_data), gr.update(value=prompt)
+        # Decode base64 image data to numpy array
+        image_bytes = base64.b64decode(image_data.split(",")[1])
+        image = Image.open(BytesIO(image_bytes))
+        image_array = np.array(image)
+        return gr.update(value=image_array), gr.update(value=prompt)
 
     entry_selector.change(
         update_selected_entry,
@@ -213,40 +222,40 @@ with gr.Blocks() as demo:
 
     def edit_entry(entry_option, new_prompt, current_dataset_name, datasets):
         if not current_dataset_name:
-            return datasets, gr.update(), gr.update(), "No dataset selected."
+            return datasets, gr.update(), gr.update(), gr.update(), "No dataset selected."
         if not entry_option or not new_prompt.strip():
-            return datasets, gr.update(), gr.update(), "Please select an entry and provide a new prompt."
+            return datasets, gr.update(), gr.update(), gr.update(), "Please select an entry and provide a new prompt."
         index = int(entry_option.split(":")[0])
         datasets[current_dataset_name][index]['prompt'] = new_prompt
         dataset = datasets[current_dataset_name]
         html_content = display_dataset_html(dataset)
         # Update entry_selector options
         entry_options = [f"{idx}: {entry['prompt'][:30]}" for idx, entry in enumerate(dataset)]
-        return datasets, gr.update(value=html_content), gr.update(choices=entry_options), f"Entry {index} updated."
+        return datasets, gr.update(value=html_content), gr.update(choices=entry_options), gr.update(value=""), f"Entry {index} updated."
 
     edit_button.click(
         edit_entry,
         inputs=[entry_selector, new_prompt_input, current_dataset_name, datasets],
-        outputs=[datasets, dataset_html, entry_selector, message_box]
+        outputs=[datasets, dataset_html, entry_selector, new_prompt_input, message_box]
     )
 
     def delete_entry(entry_option, current_dataset_name, datasets):
         if not current_dataset_name:
-            return datasets, gr.update(), gr.update(), "No dataset selected."
+            return datasets, gr.update(), gr.update(), gr.update(), gr.update(), "No dataset selected."
         if not entry_option:
-            return datasets, gr.update(), gr.update(), "Please select an entry to delete."
+            return datasets, gr.update(), gr.update(), gr.update(), gr.update(), "Please select an entry to delete."
         index = int(entry_option.split(":")[0])
         del datasets[current_dataset_name][index]
         dataset = datasets[current_dataset_name]
         html_content = display_dataset_html(dataset)
         # Update entry_selector options
         entry_options = [f"{idx}: {entry['prompt'][:30]}" for idx, entry in enumerate(dataset)]
-        return datasets, gr.update(value=html_content), gr.update(choices=entry_options), f"Entry {index} deleted."
+        return datasets, gr.update(value=html_content), gr.update(choices=entry_options), gr.update(value=None), gr.update(value=""), f"Entry {index} deleted."
 
     delete_button.click(
         delete_entry,
         inputs=[entry_selector, current_dataset_name, datasets],
-        outputs=[datasets, dataset_html, entry_selector, message_box]
+        outputs=[datasets, dataset_html, entry_selector, selected_image, selected_prompt, message_box]
     )
 
     with gr.Tab("Download Dataset"):