import streamlit as st
import json
import requests
import io 

### Inputs
x = st.text_input('Add your prompt here')

###Load variables
model_location = 'Mobius-labs/JenL_demo'


### API Request
#API_TOKEN = st.secrets["API_TOKEN"]
API_TOKEN = "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"

#API_URL = "https://api-inference.huggingface.co/models/" + model_location
API_URL = "https://lbm6i9sg7qsxwtf5.us-east-1.aws.endpoints.huggingface.cloud"

headers = {"Authorization": f"Bearer {API_TOKEN}"}

def query(payload):
    data = json.dumps(payload)
    response = requests.request("POST", API_URL, headers=headers, data=data)

    return io.BytesIO(response.content)


        
### Button
def process_api_output():
    data = query(x)

    st.image(data, output_format="JPEG")

b = st.button('Submit prompt')

if b:
    process_api_output()

#st.write(x, 'model_name is', model_name)
#st.write(x, 'model_location is', model_location)