import streamlit as st
from transformers import pipeline
import torch

# Initialize the text generation pipeline with a CPU device
generator = pipeline("text-generation", model="gpt2", device=torch.device('cpu'))

# Streamlit app layout
st.title("AI Story Generator")
st.write("Generate stories based on your custom prompt!")

# Input prompt for the user
prompt = st.text_input("Enter a story prompt:")

# Generate button to start the story generation
if st.button("Generate Story"):
    if prompt:
        try:
            # Generate story based on the prompt
            result = generator(prompt, max_length=100, num_return_sequences=1)
            
            # Display the generated story
            st.subheader("Generated Story:")
            st.write(result[0]['generated_text'])

        except Exception as e:
            # Display any errors that occur during generation
            st.error(f"An error occurred: {str(e)}")
    else:
        st.warning("Please enter a prompt to generate a story.")