app.py

import gradio as gr
import os
import zipfile
import json
from io import BytesIO
import base64
from PIL import Image
import uuid
import tempfile
import numpy as np

def save_dataset_to_zip(dataset_name, dataset):
    # Function implementation remains the same
    # ...

def load_dataset_from_zip(zip_file):
    # Function implementation remains the same
    # ...

def display_dataset_html(dataset):
    # Function implementation remains the same
    # ...

with gr.Blocks() as demo:
    gr.Markdown("<h1 style='text-align: center; margin-bottom: 20px;'>Dataset Builder</h1>")
    datasets = gr.State({})
    current_dataset_name = gr.State("")
    dataset_selector = gr.Dropdown(label="Select Dataset", interactive=True)
    entry_selector = gr.Dropdown(label="Select Entry to Edit/Delete")  # Moved outside
    dataset_html = gr.HTML()
    message_box = gr.Textbox(interactive=False, label="Message")

    with gr.Tab("Create / Upload Dataset"):
        # Create / Upload Dataset components and functions
        # ...

    def create_dataset(name, datasets):
        # Function implementation remains the same
        # ...

    create_button.click(
        create_dataset,
        inputs=[dataset_name_input, datasets],
        outputs=[dataset_selector, message_box]
    )

    def upload_dataset(zip_file, datasets):
        # Function implementation remains the same
        # ...

    upload_button.click(
        upload_dataset,
        inputs=[upload_input, datasets],
        outputs=[dataset_selector, message_box]
    )

    def select_dataset(dataset_name, datasets):
        if dataset_name in datasets:
            dataset = datasets[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 dataset_name, gr.update(value=html_content), gr.update(choices=entry_options), ""
        else:
            return "", gr.update(value="<div>Select a dataset.</div>"), gr.update(choices=[]), ""

    dataset_selector.change(
        select_dataset,
        inputs=[dataset_selector, datasets],
        outputs=[current_dataset_name, dataset_html, entry_selector, message_box]
    )

    with gr.Tab("Add Entry"):
        # Add Entry components and functions
        # ...

    def add_entry(image_data, prompt, current_dataset_name, datasets):
        if not current_dataset_name:
            return datasets, gr.update(), gr.update(), "No dataset selected."
        if image_data is None or not prompt:
            return datasets, gr.update(), gr.update(), "Please provide both an image and a prompt."
        # Convert image_data to base64
        image = Image.fromarray(image_data.astype('uint8'))
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
        img_data = f"data:image/png;base64,{img_str}"
        datasets[current_dataset_name].append({'image': img_data, 'prompt': 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 added to dataset '{current_dataset_name}'."

    add_button.click(
        add_entry,
        inputs=[image_input, prompt_input, current_dataset_name, datasets],
        outputs=[datasets, dataset_html, entry_selector, message_box]
    )

    with gr.Tab("Edit / Delete Entry"):
        with gr.Column():
            # Entry selector is referenced here, defined at the top level
            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:
            return gr.update(), gr.update()
        index = int(entry_option.split(":")[0])
        entry = datasets[current_dataset_name][index]
        image_data = entry['image']
        prompt = entry['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,
        inputs=[entry_selector, current_dataset_name, datasets],
        outputs=[selected_image, selected_prompt]
    )

    def edit_entry(entry_option, new_prompt, current_dataset_name, datasets):
        if not current_dataset_name:
            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(), 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), 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, 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(), gr.update(), gr.update(), "No dataset selected."
        if not entry_option:
            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), 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, selected_image, selected_prompt, message_box]
    )

    with gr.Tab("Download Dataset"):
        download_button = gr.Button("Download Dataset")
        download_output = gr.File(label="Download Zip")
        def download_dataset(current_dataset_name, datasets):
            if not current_dataset_name:
                return None, "No dataset selected."
            if not datasets[current_dataset_name]:
                return None, "Dataset is empty."
            zip_buffer = save_dataset_to_zip(current_dataset_name, datasets[current_dataset_name])
            # Return a dictionary with 'name' and 'data'
            return {'name': f"{current_dataset_name}.zip", 'data': zip_buffer.getvalue()}, f"Dataset '{current_dataset_name}' is ready for download."
        download_button.click(
            download_dataset,
            inputs=[current_dataset_name, datasets],
            outputs=[download_output, message_box]
        )

    # Initialize dataset_selector and entry_selector
    def initialize_components(datasets):
        return gr.update(choices=list(datasets.keys())), gr.update(choices=[])

    demo.load(
        initialize_components,
        inputs=[datasets],
        outputs=[dataset_selector, entry_selector]
    )

demo.launch()