Update CXR_LLAVA_HF.py

CXR_LLAVA_HF.py

@@ -615,11 +615,8 @@ class CXRLLAVAModel(PreTrainedModel):
             images = self.vision_tower.image_processor(image, return_tensors='pt')['pixel_values']
             images = images.to(self.device)
             input_ids = self.tokenizer_image_token(prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0)
-            if self.device == 'cuda':
-                input_ids = input_ids.cuda()
-                print('using cuda')
-            else:
-                print(f'using device {self.device}')
+            input_ids = input_ids.to(self.device)
+            # print(f'using device {self.device}')
             stopping_criteria = KeywordsStoppingCriteria(["</s>"], self.tokenizer, input_ids)
 
             image_args = {"images": images}
@@ -642,11 +639,8 @@ class CXRLLAVAModel(PreTrainedModel):
             ))
             thread.start()
             generated_text = ""
-            text_len = 0
             for new_text in streamer:
                 generated_text += new_text
-                text_len += 1
-                if text_len > 200: break
 
         return generated_text