|
from __future__ import annotations |
|
|
|
from argparse import ArgumentParser |
|
|
|
import datasets |
|
import gradio as gr |
|
import numpy as np |
|
import openai |
|
|
|
from dataset_creation.generate_txt_dataset import generate |
|
|
|
|
|
def main(openai_model: str): |
|
dataset = datasets.load_dataset("ChristophSchuhmann/improved_aesthetics_6.5plus", split="train") |
|
captions = dataset[np.random.permutation(len(dataset))]["TEXT"] |
|
index = 0 |
|
|
|
def click_random(): |
|
nonlocal index |
|
output = captions[index] |
|
index = (index + 1) % len(captions) |
|
return output |
|
|
|
def click_generate(input: str): |
|
if input == "": |
|
raise gr.Error("Input caption is missing!") |
|
edit_output = generate(openai_model, input) |
|
if edit_output is None: |
|
return "Failed :(", "Failed :(" |
|
return edit_output |
|
|
|
with gr.Blocks(css="footer {visibility: hidden}") as demo: |
|
txt_input = gr.Textbox(lines=3, label="Input Caption", interactive=True, placeholder="Type image caption here...") |
|
txt_edit = gr.Textbox(lines=1, label="GPT-3 Instruction", interactive=False) |
|
txt_output = gr.Textbox(lines=3, label="GPT3 Edited Caption", interactive=False) |
|
|
|
with gr.Row(): |
|
clear_btn = gr.Button("Clear") |
|
random_btn = gr.Button("Random Input") |
|
generate_btn = gr.Button("Generate Instruction + Edited Caption") |
|
|
|
clear_btn.click(fn=lambda: ("", "", ""), inputs=[], outputs=[txt_input, txt_edit, txt_output]) |
|
random_btn.click(fn=click_random, inputs=[], outputs=[txt_input]) |
|
generate_btn.click(fn=click_generate, inputs=[txt_input], outputs=[txt_edit, txt_output]) |
|
|
|
demo.launch(share=True) |
|
|
|
|
|
if __name__ == "__main__": |
|
parser = ArgumentParser() |
|
parser.add_argument("--openai-api-key", required=True, type=str) |
|
parser.add_argument("--openai-model", required=True, type=str) |
|
args = parser.parse_args() |
|
openai.api_key = args.openai_api_key |
|
main(args.openai_model) |
|
|
