Update chat.py

chat.py

@@ -131,7 +131,7 @@ class Chat:
         keywords = [stop_str]
         stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
         streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15)
-        
+
         max_new_tokens = min(max_new_tokens, max_context_length - input_ids.shape[-1] - num_image_tokens - num_seg_tokens - num_depth_tokens)
 
         if max_new_tokens < 1:
@@ -159,30 +159,30 @@ class Chat:
             yield json.dumps({"text": generated_text, "error_code": 0}).encode()
 
     def generate_stream_gate(self, params):
-        # try:
-        for x in self.generate_stream(params):
-            yield x
-        # except ValueError as e:
-        #     print("Caught ValueError:", e)
-        #     ret = {
-        #         "text": server_error_msg,
-        #         "error_code": 1,
-        #     }
-        #     yield json.dumps(ret).encode()
-        # 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()
-        # except Exception as e:
-        #     print("Caught Unknown Error", e)
-        #     ret = {
-        #         "text": server_error_msg,
-        #         "error_code": 1,
-        #     }
-        #     yield json.dumps(ret).encode()
+        try:
+            for x in self.generate_stream(params):
+                yield x
+        except ValueError as e:
+            print("Caught ValueError:", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode()
+        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()
+        except Exception as e:
+            print("Caught Unknown Error", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode()
 
 
 if __name__ == "__main__":