Updated app.py

app.py

@@ -1,4 +1,5 @@
 import streamlit as st
+import cohere
 from PIL import Image
 from transformers import BlipProcessor, BlipForConditionalGeneration, AutoTokenizer
 import itertools
@@ -9,6 +10,10 @@ import torch
 import numpy as np
 nltk.download('stopwords')
 
+COHERE_API_KEY = os.getenv('COHERE_API_KEY')
+co_client = cohere.Client(COHERE_API_KEY)
+
+
 processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
 model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
 
@@ -27,15 +32,24 @@ def genrate_caption(image_file):
 st.title("Image Caption and HashTag Recommender")
 image_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
 
+def creative_caption(text):
+    return co_client.generate(prompt=f"Write some trendy instagram captions for the following prompt - {text}").generations[0].text
+
+
+def caption_hashtags(text):
+    return co_client.generate(prompt=f"Write some trendy instagram hashtags for the following prompt - {text}").generations[0].text
+
 if image_file is not None:
     try:
         caption = genrate_caption(image_file)
+        caption_text = creative_caption(caption)
+        hashtags = caption_hashtags(caption)
         if len(caption) > 0:
             st.write(f"Caption : {caption}")
+            st.write(f"Creative Caption : {caption_text}")
+            st.write(f"Creative hashtags : {hashtags}")
         
         else:
             st.write("No caption found for this image.")
     except Exception as e:
-        st.write(f"Error: {e}")
-
-    
+        st.write(f"Error: {e}")