import requests
import streamlit as st

ACCESS_TOKEN = st.secrets["ACCESS_TOKEN"]

API_for_caption = "https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1"
API_for_semantic = "https://api-inference.huggingface.co/models/Salesforce/blip-image-captioning-large"
headers = {"Authorization": ACCESS_TOKEN}


def semantic_generator(filename):
    response = requests.post(API_for_semantic, headers=headers, data=filename)
    semantic = response.json()
    print(f"semantic: {semantic[0]['generated_text']}")
    return semantic[0]['generated_text']


def caption_generator(semantic):
    prompt = {
	"inputs": f"write a facebook caption for following text semantic for my image.Do not include the semantic in the answer.just give the answer.make sure to add emojis and hashtags"
              f"semantic:{semantic} \
                Answer:",}
    response = requests.post(API_for_caption, headers=headers, json=prompt)
    caption_raw = response.json()
    caption = caption_raw[0]['generated_text'].split('Answer:')
    print(f"caption: {caption}")
    return caption[1]
	

st.title('Caption Generator')
image = st.file_uploader(label='Add an image',type=['jpg','png','webp'])

if image:
    col1,col2 = st.columns(2)
     
    with col1:
        st.image(image)

    with col2:
        st.subheader('Semantic:')
        with st.spinner('Generating semantic'):
            semantic = semantic_generator(image)
            st.markdown(semantic)

        st.subheader('Caption:')
        with st.spinner('Generating caption'):
            caption = caption_generator(semantic)
            st.markdown(caption)