Update app.py

app.py

@@ -14,20 +14,14 @@ except ImportError:
     st.success("PyTorch has been successfully installed!")
     import torch
 
-# Install Optimum library
-try:
-    import optimum
-except ImportError:
-    st.warning("Optimum library is not installed. Installing optimum...")
-    import subprocess
-    subprocess.run(["pip", "install", "optimum"])
-    st.success("Optimum library has been successfully installed!")
-    import optimum
-
 # Load the image captioning model
 caption_model = pipeline("image-to-text", model="unography/blip-large-long-cap")
 
-story_generator = pipeline("text-generation", model="TheBloke/storytime-13B-GPTQ")
+#story_generator = pipeline("text-generation", model="distilbert/distilgpt2")
+
+#story_generator = pipeline("text-generation", model="isarth/distill_gpt2_story_generator")
+
+story_generator = pipeline("text-generation", model="mosaicml/mpt-7b-storywriter", trust_remote_code=True)
 
 def generate_caption(image):
     # Generate the caption for the uploaded image
@@ -36,7 +30,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 story inspired by the image of {caption}. Here's the story:\n\n"
+    prompt = f"Write a short, simple children's fairytale story with reference of 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