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 10K 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
from transformers import load_tool

prompt_assistant = load_tool(
    "llamaindex/text2image_prompt_assistant",
    openai_api_key="your_api_key",
    model_name='text-davinci-003',
    temperature=0.3,  # increase or decrease this to control variation
    verbose=True
)

agent = OpenAiAgent(model="text-davinci-003", api_key="your_api_key")
agent.run("Draw me a picture a river and some trees.")
```
"""
)

st.subheader("Try out the tool below!")
st.markdown("Note that the first run may be slow as it downloads files needed inference. Furthermore, the text-2-image model itself will be running on CPU. After the first run, it should be faster!")

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 image'):
    with st.spinner('Running...'):
        tool = Text2ImagePromptAssistant(openai_api_key=api_key, model_name=model_name, temperature=temperature)
        response = tool(initial_input)
        st.image(response)