Update app.py

app.py

@@ -18,6 +18,8 @@ from PIL import Image
 import torch
 import spaces
 import subprocess
+import os
+
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 torch.set_default_device('cuda')
@@ -108,29 +110,43 @@ tokenizer = AutoTokenizer.from_pretrained("5CD-AI/Vintern-3B-beta", trust_remote
 
 @spaces.GPU
 def chat(message, history):
-    print(history)
-    print(message)
+    print("history",history)
+    print("message",message)
+        
     if len(history) == 0 or len(message["files"]) != 0:
         test_image = message["files"][0]["path"]
-    else:
+        pixel_values = load_image(test_image, max_num=6).to(torch.bfloat16).cuda()
+    elif len(history) == 0 or len(message["files"]) == 0:
+        pixel_values = None     
+    elif history[0][0][0] is not None and os.path.isfile(history[0][0][0]):
         test_image = history[0][0][0]
+        pixel_values = load_image(test_image, max_num=6).to(torch.bfloat16).cuda()
+    else:
+        pixel_values = None 
         
-    pixel_values = load_image(test_image, max_num=6).to(torch.bfloat16).cuda()
-    generation_config = dict(max_new_tokens= 512, do_sample=False, num_beams = 3, repetition_penalty=4.0)
-    
     
+    generation_config = dict(max_new_tokens= 512, do_sample=False, num_beams = 3, repetition_penalty=2.0)
     
     if len(history) == 0:
-        question = '<image>\n'+message["text"]
+        if pixel_values is not None:
+            question = '<image>\n'+message["text"]
+        else:
+            question = message["text"]
         response, conv_history = model.chat(tokenizer, pixel_values, question, generation_config, history=None, return_history=True)
     else:
+        ### get the last image
         conv_history = []
         for chat_pair in history:
             if chat_pair[1] is not None:
                 if len(conv_history) == 0 and len(message["files"]) == 0:
                     chat_pair[0] = '<image>\n' + chat_pair[0]
                 conv_history.append(tuple(chat_pair))
-        print(conv_history)
+        ### not have image
+        if len(conv_history) == 0:
+            for chat_pair in history:
+                conv_history.append(tuple(chat_pair))
+                
+        print("conv_history",conv_history)
         if len(message["files"]) != 0:
             question = '<image>\n'+message["text"]
         else: