Update app.py

app.py

@@ -17,8 +17,8 @@ except ImportError:
 # Load the image captioning model
 caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
 
-# Load the DeepSeek model for story generation
-story_generator = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1", trust_remote_code=True)
+# Load the GPT-2 model for story generation
+story_generator = pipeline("text-generation", model="gpt2")
 
 def generate_caption(image):
     # Generate the caption for the uploaded image
@@ -26,10 +26,9 @@ def generate_caption(image):
     return caption
 
 def generate_story(caption):
-    # Generate the story based on the caption using the DeepSeek model
+    # Generate the story based on the caption using the GPT-2 model
     prompt = f"Imagine you are a storyteller for young children. Based on the image described as '{caption}', create a short and interesting story for children aged 3-10. Keep it positive and happy in tone."
-    messages = [{"role": "user", "content": prompt}]
-    story = story_generator(messages)[0]["generated_text"]
+    story = story_generator(prompt, max_length=200, num_return_sequences=1)[0]["generated_text"]
     return story
 
 def convert_to_audio(story):
@@ -58,7 +57,7 @@ def main():
         st.subheader("Generated Caption:")
         st.write(caption)
         
-        # Generate the story based on the caption using the DeepSeek model
+        # Generate the story based on the caption using the GPT-2 model
         story = generate_story(caption)
         st.subheader("Generated Story:")
         st.write(story)