import streamlit as st
from prompt_assistant import Text2ImagePromptAssistant

st.title('LlamaIndex Text2ImagePromptAssistant')

st.markdown('This is a tool for HF agents, that assists with writing text-to-image prompts with the power of [LlamaIndex](https://gpt-index.readthedocs.io/en/latest/index.html). By indexing 50K random prompts from [DiffusionDB](https://huggingface.co/datasets/poloclub/diffusiondb), LlamaIndex is able to suggest text-to-image prompts that will hopefully lead to more beautiful results. With a HF Agent, the agent will call the prompt assistant tool before generating an image from text.')

st.markdown('Example usage is below:')

st.markdown("""
```python
TODO
```
"""
)

st.subheader("Try out the tool below!")

api_key = st.text_input('OpenAI API key')

model_name = st.selectbox('Model Name', options=['text-davinci-003', 'gpt-3.5-turbo', 'gpt-4'])

temperature = st.slider('Temperature', min_value=0.0, max_value=1.0, step=0.05, value=0.3)

initial_input = st.text_input('Initial text-to-image prompt', placeholder='Draw me an image of a happy dog.')

if st.button('Generate optimized prompt'):
    with st.spinner('Running...'):
        tool = Text2ImagePromptAssistant(openai_api_key=api_key, model_name=model_name, temperature=temperature)
        response = tool(initial_input)
        st.markdown(response)