import gradio as gr

from askui_ml_helper.utils.pta_text import PtaTextInference

def main():
    inference = PtaTextInference("model/pta-text-v0.1.pt")
    # Gradio Interface
    iface = gr.Interface(
        fn=inference.process_image_and_draw_circle,
        inputs=[gr.Image(type="pil", label = "Upload Image"), 
                gr.Textbox(label = "Prompt", placeholder='Enter a prompt to find in the image', value='Click on the text "Login"')],
        outputs=gr.Image(type="pil"),
        title="PTA-Text: A **Text-only** Click Model",
        description="Based on PTA-text ([Huggingface model card](https://huggingface.co/AskUI/pta-text-0.1) ).PTA is a small part of our bigger model - PTA-1 (Prompt-To-Automation). \nInstructions: Upload a UI image and enter a text finding prompt to see the model's prediction. \nLimitations: Currently, limited to single screen images."
    )

    iface.launch(server_name="0.0.0.0")


if __name__ == "__main__":
    main()