import streamlit as st
from transformers import pipeline, set_seed

# Create a text generation pipeline with the "google/flan-t5-base" model
text_generator = pipeline("text-generation", model="google/flan-t5-base")

st.title("Story Generator")

# User inputs
title = st.text_input("Title of the Story:")
character = st.text_input("Character (Hero) Name:")
era = st.selectbox("Era of the Story:", ["Medieval", "Modern", "Future"])
genre = st.selectbox("Genre:", ["Action", "Novelistic", "Love"])
ending = st.selectbox("Ending:", ["Happy", "Sad"])

if st.button("Generate Story"):
    if title and character:
        # Prepare the prompt based on user inputs
        prompt = f"In the {era} era, there was a character named {character}. This is the {genre} story of {character}, a {genre} story with a {ending} ending."
        
        # Generate the story
        set_seed(42)  # Set a seed for reproducibility
        story = text_generator(prompt, max_length=500, do_sample=True)[0]["generated_text"]
        
        # Set the title
        full_story = f"# {title}\n\n{story}"
        
        st.markdown(full_story)
    else:
        st.warning("Please provide a title and character name.")