import gradio as gr

from src.v2_for_hf import generate_images
from src.v2_for_hf import NUM_GEN

iface = gr.Interface(
    fn=generate_images,
    inputs=[
        gr.Textbox(label="OpenAI API Key"),
        gr.Image(label="Input Image", type="filepath"),
        gr.Textbox(label="Mistaken Class"),
        gr.Textbox(label="Ground Truth Class")
    ],
    outputs=[
        gr.Image(label="Generated Image of Ground Truth Class") for i in range(NUM_GEN)
    ],
    title="visual-data-aug",
)

if __name__ == "__main__":
    iface.launch(share=True)