import streamlit as st
from transformers import pipeline
from huggingface_hub import InferenceClient
from PIL import Image
import os


def initialize():
    if 'initialized' not in st.session_state: 
        print("Initializing...")
        st.session_state['initialized'] = True
        st.session_state['api_key'] = os.getenv("HUGGINGFACE_TOKEN")
        st.session_state['client'] = InferenceClient(api_key=st.session_state['api_key'])


def main():
    initialize()
    st.header("Character Captions")
    st.write("Have a character caption any image you upload!")
    character = st.selectbox("Choose a character", ["artist", "elmo", "unintelligible", "goku"])

    uploaded_img = st.file_uploader("Upload an image here")

    if uploaded_img is not None:
        image = Image.open(uploaded_img)
        st.image(image)

        
        image_captioner = pipeline("image-to-text", model="Salesforce/blip-image-captioning-large")
        response = image_captioner(image)
        caption = response[0]['generated_text']


        
        character_prompts = {
            "artist": f"Describe this caption like you're a artist: {caption}.",
            "elmo": f"Describe this caption like you're elmo: {caption}.",
            "unintelligible": f"Describe this caption in a way that makes no sense: {caption}.",
            "goku": f"Describe this caption like you're goku: {caption}."
        }

        prompt = character_prompts[character]
        messages = [
            { "role": "user", "content": prompt }
        ]

        
        stream = st.session_state['client'].chat.completions.create(
            model="meta-llama/Llama-3.2-3B-Instruct", 
            messages=messages, 
            max_tokens=500,
            stream=True
        )

        response = ''
        for chunk in stream:
            response += chunk.choices[0].delta.content
        
        st.write(response)



if __name__ == '__main__':
    main()