Spaces:
Running
on
T4
Running
on
T4
| 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...") # fmt: skip | |
| 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) | |