Update app.py

app.py

@@ -6,7 +6,7 @@ from threading import Thread
 # import time
 import cv2
 
-
+import datetime
 # import copy
 import torch
 
@@ -34,8 +34,6 @@ from llava.mm_utils import (
 
 from serve_constants import html_header
 
-from PIL import Image
-
 import requests
 from PIL import Image
 from io import BytesIO
@@ -46,6 +44,9 @@ import gradio_client
 import subprocess
 import sys
 
+external_log_dir = "./logs"
+LOGDIR = external_log_dir
+
 def install_gradio_4_35_0():
     current_version = gr.__version__
     if current_version != "4.35.0":
@@ -64,6 +65,11 @@ import gradio_client
 print(f"Gradio version: {gr.__version__}")
 print(f"Gradio-client version: {gradio_client.__version__}")
 
+def get_conv_log_filename():
+    t = datetime.datetime.now()
+    name = os.path.join(LOGDIR, f"{t.year}-{t.month:02d}-{t.day:02d}-user_conv.json")
+    return name
+
 class InferenceDemo(object):
     def __init__(
         self, args, model_path, tokenizer, model, image_processor, context_len
@@ -113,6 +119,16 @@ def is_valid_video_filename(name):
     else:
         return False
 
+def is_valid_image_filename(name):
+    image_extensions = ["jpg", "jpeg", "png", "bmp", "gif", "tiff", "webp", "heic", "heif", "jfif", "svg", "eps", "raw"] 
+
+    ext = name.split(".")[-1].lower()
+
+    if ext in image_extensions:
+        return True
+    else:
+        return False
+
 
 def sample_frames(video_file, num_frames):
     video = cv2.VideoCapture(video_file)
@@ -193,9 +209,14 @@ def bot(history):
         if type(message[0]) is tuple:
             images_this_term.append(message[0][0])
             if is_valid_video_filename(message[0][0]):
+                # 不接受视频
+                raise ValueError("Video is not supported")
                 num_new_images += our_chatbot.num_frames
-            else:
+            elif is_valid_image_filename(message[0][0]):
+                print("#### Load image from local file",message[0][0])
                 num_new_images += 1
+            else:
+                raise ValueError("Invalid image file")
         else:
             num_new_images = 0
 
@@ -209,8 +230,11 @@ def bot(history):
     for f in images_this_term:
         if is_valid_video_filename(f):
             image_list += sample_frames(f, our_chatbot.num_frames)
-        else:
+        elif is_valid_image_filename(f):
             image_list.append(load_image(f))
+        else:
+            raise ValueError("Invalid image file")
+    
     image_tensor = [
         our_chatbot.image_processor.preprocess(f, return_tensors="pt")["pixel_values"][
             0
@@ -219,6 +243,24 @@ def bot(history):
         .to(our_chatbot.model.device)
         for f in image_list
     ]
+    all_image_hash = []
+    for image_path in image_list:
+        with open(image_path, "rb") as image_file:
+            image_data = image_file.read()
+            image_hash = hashlib.md5(image_data).hexdigest()
+            all_image_hash.append(image_hash)
+            image = PIL.Image.open(image_path).convert("RGB")
+            all_images.append(image)
+            t = datetime.datetime.now()
+            filename = os.path.join(
+                LOGDIR,
+                "serve_images",
+                f"{t.year}-{t.month:02d}-{t.day:02d}",
+                f"{image_hash}.jpg",
+            )
+            if not os.path.isfile(filename):
+                os.makedirs(os.path.dirname(filename), exist_ok=True)
+                image.save(filename)
 
     image_tensor = torch.stack(image_tensor)
     image_token = DEFAULT_IMAGE_TOKEN * num_new_images
@@ -280,7 +322,19 @@ def bot(history):
     our_chatbot.conversation.messages[-1][-1] = outputs
 
     history[-1] = [text, outputs]
-
+    print("#### history",history)
+    
+    with open(get_conv_log_filename(), "a") as fout:
+        data = {
+            "tstamp": round(finish_tstamp, 4),
+            "type": "chat",
+            "model": "Pangea-7b",
+            "start": round(start_tstamp, 4),
+            "finish": round(start_tstamp, 4),
+            "state": history,
+            "images": all_image_hash,
+        }
+        fout.write(json.dumps(data) + "\n")
     return history
     # generate_kwargs = dict(
     #     inputs=input_ids,
@@ -345,7 +399,7 @@ with gr.Blocks(
     
     with gr.Column():
         with gr.Row():
-            chatbot = gr.Chatbot([], elem_id="chatbot", bubble_full_width=False)
+            chatbot = gr.Chatbot([], elem_id="Pangea", bubble_full_width=False, height=750)
 
         with gr.Row():
             upvote_btn = gr.Button(value="👍  Upvote", interactive=True)
@@ -560,8 +614,8 @@ if __name__ == "__main__":
     argparser.add_argument("--model-base", type=str, default=None)
     argparser.add_argument("--num-gpus", type=int, default=1)
     argparser.add_argument("--conv-mode", type=str, default=None)
-    argparser.add_argument("--temperature", type=float, default=0.2)
-    argparser.add_argument("--max-new-tokens", type=int, default=512)
+    argparser.add_argument("--temperature", type=float, default=0.7)
+    argparser.add_argument("--max-new-tokens", type=int, default=4096)
     argparser.add_argument("--num_frames", type=int, default=16)
     argparser.add_argument("--load-8bit", action="store_true")
     argparser.add_argument("--load-4bit", action="store_true")