Update app.py

app.py

@@ -1,31 +1,30 @@
 import streamlit as st
-import io
 from transformers import pipeline, set_seed
 
-# Create a text2text-generation pipeline with the "google/flan-t5-base" model
-pipe = pipeline("text2text-generation", model="google/flan-t5-base")
+# Create a text generation pipeline with the "gpt2" model
+text_generator = pipeline("text-generation", model="gpt2")
 
-# Initialize the EasyOCR reader for text extraction from images
-# ocr_reader = easyocr.Reader(['en'])
+st.title("Story Generator")
 
-# Remove the text extraction code
-
-# Create input fields for the title, character, era, genre, and ending
-title = st.text_input("Title:")
-character = st.text_input("Character (hero):")
-era = st.text_input("Era of story:")
+# 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"])
 
-# Generate a story prompt based on the user input
-story_prompt = f"Once upon a time, in the {era}, there lived a {character} named {character}. {character} was a {genre} hero/heroine, and one day, they went on a quest to..."
-
-# Set a seed for reproducibility
-set_seed(42)
-
-# Generate the story
-story = pipe(story_prompt, max_length=300, do_sample=True)[0]["generated_text"]
-
-# Display the story to the user
-st.markdown("**Story:**")
-st.markdown(story)
+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=300, 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.")