Update chat.py

chat.py

@@ -11,9 +11,10 @@ from vcoder_llava.mm_utils import process_images, load_image_from_base64, tokeni
 from vcoder_llava.constants import (
     IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN,
     SEG_TOKEN_INDEX, DEFAULT_SEG_TOKEN,
-    DEPTH_TOKEN_INDEX, DEFAULT_DEPTH_TOKEN
+    DEPTH_TOKEN_INDEX, DEFAULT_DEPTH_TOKEN,
 )
 from transformers import TextIteratorStreamer
+from threading import Thread
 
 class Chat:
     def __init__(self, model_path, model_base, model_name,
@@ -35,7 +36,7 @@ class Chat:
             model_path, model_base, self.model_name, load_8bit, load_4bit, device=self.device)
         self.is_multimodal = 'llava' in self.model_name.lower()
         self.is_seg = "vcoder" in self.model_name.lower()
-        self.is_depth = False
+        self.is_depth = "ds" in self.model_name.lower()
 
     @torch.inference_mode()
     def generate_stream(self, params):
@@ -167,21 +168,21 @@ class Chat:
                 "text": server_error_msg,
                 "error_code": 1,
             }
-            yield json.dumps(ret).encode()
+            yield json.dumps(ret).encode() + b"\0"
         except torch.cuda.CudaError as e:
             print("Caught torch.cuda.CudaError:", e)
             ret = {
                 "text": server_error_msg,
                 "error_code": 1,
             }
-            yield json.dumps(ret).encode()
+            yield json.dumps(ret).encode() + b"\0"
         except Exception as e:
             print("Caught Unknown Error", e)
             ret = {
                 "text": server_error_msg,
                 "error_code": 1,
             }
-            yield json.dumps(ret).encode()
+            yield json.dumps(ret).encode() + b"\0"
 
 
 if __name__ == "__main__":