Update app.py

app.py

@@ -14,13 +14,19 @@ 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="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):