import streamlit as st
from huggingface_hub import HfApi, ModelFilter, InferenceClient

if 'models' not in st.session_state:

    hf_api = HfApi()
    st.session_state.models = sorted(list(hf_api.list_models(
        filter=ModelFilter(
            task="text-to-image",
            library="diffusers",
        )
    )), key=lambda x: x.downloads, reverse=True)

def format_model(model):
    return f"{model.modelId} (⬇️ {model.downloads}, ❤️ {model.likes})"


st.title("Text to Image Model testing")
st.session_state.token = st.text_input("Enter your API token", type="password")
user_input = st.text_input("Your prompt")
model_name = st.selectbox("Select", st.session_state.models, format_func=format_model)

if model_name and user_input and st.session_state.token:
    model_name = model_name.modelId
    client = InferenceClient(model_name, token=st.session_state.token)
    pred = client.text_to_image(user_input)
    st.image(pred)