Update app.py

app.py

@@ -14,20 +14,13 @@ except ImportError:
     st.success("PyTorch has been successfully installed!")
     import torch
 
-# Install Einops library
-try:
-    import einops
-except ImportError:
-    st.warning("Einops library is not installed. Installing einops...")
-    import subprocess
-    subprocess.run(["pip", "install", "einops"])
-    st.success("Einops library has been successfully installed!")
-    import einops
-
 # Load the image captioning model
 caption_model = pipeline("image-to-text", model="unography/blip-large-long-cap")
 
-story_generator = pipeline("text-generation", model="ajibawa-2023/Young-Children-Storyteller-Mistral-7B")
+#story_generator = pipeline("text-generation", model="distilbert/distilgpt2")
+
+story_generator = pipeline("text-generation", model="isarth/distill_gpt2_story_generator")
+
 
 def generate_caption(image):
     # Generate the caption for the uploaded image
@@ -36,7 +29,7 @@ def generate_caption(image):
 
 def generate_story(caption):
     # Generate the story based on the caption using the GPT-2 model
-    prompt = f"Write a short, simple children's fairytale story with reference of the image of {caption}. Here's the story:\n\n"
+    prompt = f"Write a short, simple children's story inspired by the image of {caption}. Here's the story:\n\n"
     story = story_generator(prompt, max_length=500, num_return_sequences=1)[0]["generated_text"]
     
     # Extract the story text from the generated output