debug

model_worker.py

@@ -52,12 +52,12 @@ class ModelWorker:
                 torch_device='cpu',
                 device_map="cpu",
             )
-        self.model.to("cuda:0")
+        self.model.to('cuda')
 
     @spaces.GPU
     def generate_stream(self, params):
         tokenizer, model, image_processor = self.tokenizer, self.model, self.image_processor
-        logger.info(f'Model devices: {self.model.device}')
+        logger.info(f'Model devices: {model.device}')
 
         prompt = params["prompt"]
         ori_prompt = prompt
@@ -70,17 +70,18 @@ class ModelWorker:
 
                 images = [load_image_from_base64(image) for image in images]
                 images = process_images(images, image_processor, model.config)
+                logger.info(f'Images: {images.shape}')
 
                 if type(images) is list:
-                    images = [image.to(self.model.device, dtype=torch.float16) for image in images]
+                    images = [image.to(model.device, dtype=torch.float16) for image in images]
                 else:
-                    images = images.to(self.model.device, dtype=torch.float16)
+                    images = images.to(model.device, dtype=torch.float16)
 
                 if self.load_bf16:
                     images = images.to(dtype=torch.bfloat16)
 
                 replace_token = DEFAULT_IMAGE_TOKEN
-                if getattr(self.model.config, 'mm_use_im_start_end', False):
+                if getattr(model.config, 'mm_use_im_start_end', False):
                     replace_token = DEFAULT_IM_START_TOKEN + replace_token + DEFAULT_IM_END_TOKEN
                 prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token)
 
@@ -99,15 +100,15 @@ 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).cuda()
+        input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
         keywords = [stop_str]
         stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
-        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15)
+        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=None)
 
         max_new_tokens = min(max_new_tokens, max_context_length - input_ids.shape[-1] - num_image_tokens)
 
         if max_new_tokens < 1:
-            yield json.dumps({"text": ori_prompt + "Exceeds max token length. Please start a new conversation, thanks.", "error_code": 0}).encode() + b"\0"
+            yield json.dumps({"text": ori_prompt + "Exceeds max token length. Please start a new conversation, thanks.", "error_code": 0}).encode()
             return
 
         thread = Thread(target=model.generate, kwargs=dict(
@@ -128,33 +129,11 @@ class ModelWorker:
             generated_text += new_text
             if generated_text.endswith(stop_str):
                 generated_text = generated_text[:-len(stop_str)]
-            yield json.dumps({"text": generated_text, "error_code": 0}).encode() + b"\0"
+            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() + 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"
+        for x in self.generate_stream(params):
+            yield x
 
 def release_model_semaphore(fn=None):
     model_semaphore.release()

requirements.txt

@@ -11,4 +11,5 @@ einops==0.6.1
 einops-exts==0.0.4
 timm==0.6.13
 httpx==0.24.0
+numpy==1.26.4
 scipy