import os
from pathlib import Path
import streamlit as st
from transformers import pipeline
from dotenv import load_dotenv
from langchain import PromptTemplate, HuggingFaceHub, LLMChain

if Path(".env").is_file():
    load_dotenv(".env")
st.set_page_config(layout="wide")
HF_TOKEN = os.getenv("HF_TOKEN")

def img2Text(url):
    image_to_text = pipeline("image-to-text", model="Salesforce/blip-image-captioning-large")
    text = image_to_text(url)
    st.subheader("Caption :")
    st.subheader(text)
    print(text)
    return text
img2Text("photo.jpg")

#llm
def generate_story(scenario):
    template = """
    You are a story teller;
    You can generate a short story based on a simple narrative, the story should be no momre than 20 words;
    CONTEXT: {scenario}
    STORY:
    """

    prompt = PromptTemplate(template=template,input_variables=["scenario"])
    llm_chain = LLMChain(prompt=prompt, 
                     llm=HuggingFaceHub(repo_id="google/flan-t5-xl", 
                                        model_kwargs={"temperature":0, 
                                                      "max_length":64}))
    story =llm_chain.run(scenario)
    st.subheader("Story :")
    st.subheader(story)
    print(story)
    return story

scenario = img2Text("photo.jpg")
story = generate_story(scenario)