Update bunny/serve/model_worker.py

bunny/serve/model_worker.py

@@ -154,10 +154,8 @@ class ModelWorker:
         stop_str = params.get("stop", None)
         do_sample = True if temperature > 0.001 else False
 
-        
         input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(
             self.device)
-
         keywords = [stop_str]
         stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
         streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15)
@@ -173,7 +171,6 @@ class ModelWorker:
 
         model = model.to('cuda')
 
-
         thread = Thread(target=model.generate, kwargs=dict(
             inputs=input_ids,
             do_sample=do_sample,
@@ -199,33 +196,33 @@ class ModelWorker:
             yield json.dumps({"text": generated_text, "error_code": 0}).encode() + b"\0"
 
     def generate_stream_gate(self, params):
-        # for x in self.generate_stream(params):
-        #     yield x
-
-        try:
-            for x in self.generate_stream(params):
-                yield x
-        except ValueError as e:
-            logger.info("Caught ValueError:", e)
-            ret = {
-                "text": server_error_msg,
-                "error_code": 1,
-            }
-            yield json.dumps(ret).encode() + b"\0"
-        except torch.cuda.CudaError as e:
-            logger.info("Caught torch.cuda.CudaError:", e)
-            ret = {
-                "text": server_error_msg,
-                "error_code": 1,
-            }
-            yield json.dumps(ret).encode() + b"\0"
-        except Exception as e:
-            logger.info("Caught Unknown Error", e)
-            ret = {
-                "text": server_error_msg,
-                "error_code": 1,
-            }
-            yield json.dumps(ret).encode() + b"\0"
+        for x in self.generate_stream(params):
+            yield x
+
+        # 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() + 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() + b"\0"
+        # except Exception as e:
+        #     print("Caught Unknown Error", e)
+        #     ret = {
+        #         "text": server_error_msg,
+        #         "error_code": 1,
+        #     }
+        #     yield json.dumps(ret).encode() + b"\0"
 
 
 app = FastAPI()