Update app.py

app.py

@@ -1,30 +1,25 @@
 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'))
+# Initialize the Hugging Face pipeline for text generation
+generator = pipeline("text-generation", model="gpt2")
 
 # Streamlit app layout
-st.title("AI Story Generator")
-st.write("Generate stories based on your custom prompt!")
+st.title("AI Story Generator - Creativista ")
 
-# Input prompt for the user
-prompt = st.text_input("Enter a story prompt:")
+# Prompt input from the user
+prompt = st.text_input("Enter a prompt:", "")
 
-# Generate button to start the story generation
+# Slider to control the length of the generated story
+max_length = st.slider("Select the maximum length of the story (in words):", 
+                       min_value=50, max_value=500, value=100, step=10)
+
+# Button to trigger 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)}")
+        # Generate the story based on the prompt
+        result = generator(prompt, max_length=max_length, num_return_sequences=1)
+        # Display the generated story
+        st.write(result[0]['generated_text'])
     else:
-        st.warning("Please enter a prompt to generate a story.")
+        st.write("Please enter a prompt to generate a story.")