import gradio as gr
from datasets import load_dataset



startimg = 0

def get_dataset_forward():
    global startimg
    final = []
    dataset = load_dataset("ShoukanLabs/OpenNiji-Dataset-Aesthetic-Finetune", split=f"train[{startimg}:{startimg + 50}]")
    for idx in dataset:
        url = idx["url"]
        prompt = idx["prompt"]
        style = idx["style"]
        score = idx["score"]
        final.append((url, f"{prompt}\n\n Style: {style}\n\n Aesthetic: {score}"))
    startimg += 50
    return final

def get_dataset_back():
    global startimg
    final = []
    startimg -= 50
    dataset = load_dataset("ShoukanLabs/OpenNiji-Dataset", split=f"train[{startimg}:{startimg + 50}]")
    for idx in dataset:
        url = idx["url"]
        prompt = idx["prompt"]
        style = idx["style"]
        score = idx["score"]
        final.append((url, f"{prompt}\n\n Style: {style}\n\n Aesthetic: {score}"))
    return final

with gr.Blocks() as demo:
    with gr.Column():
        with gr.Row():
            back = gr.Button("<").style()
            forward = gr.Button(">").style()
        gallery = gr.Gallery(
            label="Showing 50 images", show_label=True, elem_id="gallery"
        ).style(object_fit="contain", columns=[10], height="auto")

        back.click(get_dataset_back, None, gallery)
        forward.click(get_dataset_forward, None, gallery)

if __name__ == "__main__":
    demo.launch()