Update app.py

app.py

@@ -14,13 +14,19 @@ 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="distilbert/distilgpt2")
-
-#story_generator = pipeline("text-generation", model="isarth/distill_gpt2_story_generator")
-
 story_generator = pipeline("text-generation", model="TheBloke/storytime-13B-GPTQ")
 
 def generate_caption(image):