Update application/chat_inference.py

application/chat_inference.py

@@ -1,10 +1,12 @@
 from application.utils.chat_completion_api import ChatCompletionAPI
 from config import Response,pipeline_dict,convs_dict
 import os
+from application.utils.image_captioning import ImageCaptioning
 
 class ChatInference:
     def __init__(self):
         self.chatCompletionAPI = ChatCompletionAPI()
+        self.image_captioning = ImageCaptioning()
 
     def validate(self,data,user):
         try:
@@ -30,12 +32,19 @@ class ChatInference:
             top_p = max(0.1, min(data.get('top_p', 0.9), 1))
             system = data.get('system_prompt','You are a helpful and harmless AI assistant. You are xylaria made by sk md saad amin. You should think step-by-step')
             convId = data['convId']
-            
+            image = data.get('image')
+
             if(len(convs_dict[user][convId]['messages'])==1):
                 if system:
                     convs_dict[user][convId]['messages'].insert(0,{"role":"system", "content": system})
                 convs_dict[user]['metadata'].insert(0,{"convId": convId, "title": prompt[:23]})
                 convs_dict[user][convId]['title'] = prompt[:30]
+
+            if image:
+                caption = self.image_captioning.generate_caption(image)
+                prompt = f"{caption}\n\n{prompt}"
+
+
             if(pipeline[model]['type'] == 'image-text-to-text'):
                 convs_dict[user][convId]['messages'].append({"role": "user", "content": [{"type":"text","text":prompt}]})
             else: