Update handler.py

handler.py

@@ -1,12 +1,12 @@
 import json
 import torch
-from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
+from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, pipeline
 from qwen_vl_utils import process_vision_info
 
 
 class EndpointHandler:
     def __init__(self, model_dir):
-        # Load the model and processor for Qwen2-VL
+        # Initialize the model and processor for Visual Question Answering (VQA)
         self.model = Qwen2VLForConditionalGeneration.from_pretrained(
             model_dir,
             torch_dtype=torch.float16,  # FP16 for memory efficiency
@@ -16,6 +16,13 @@ class EndpointHandler:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.model.eval()
 
+        # Initialize the VQA pipeline
+        self.vqa_pipeline = pipeline(
+            task="visual-question-answering",
+            model=self.model,
+            device=0 if torch.cuda.is_available() else -1
+        )
+
     def preprocess(self, request_data):
         # Parse messages, extract video and text inputs
         messages = request_data.get('messages')
@@ -41,42 +48,31 @@ class EndpointHandler:
         return inputs.to(self.device)
 
     def inference(self, inputs):
-        # Run inference on the model
+        # Use the VQA pipeline for inference
         with torch.no_grad():
-            generated_ids = self.model.generate(
-                **inputs, 
-                max_new_tokens=128,  # Limit the output length
-                num_beams=1,  # Reduce memory usage
-                max_batch_size=1  # Process one batch at a time
+            result = self.vqa_pipeline(
+                images=inputs["images"] if "images" in inputs else inputs["videos"],
+                question=inputs["text"]
             )
 
-        # Trim generated outputs to remove input tokens
-        generated_ids_trimmed = [
-            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-        ]
-
-        return generated_ids_trimmed
+        return result
 
     def postprocess(self, inference_output):
-        # Decode generated output into human-readable text
-        output_text = self.processor.batch_decode(
-            inference_output, skip_special_tokens=True, clean_up_tokenization_spaces=False
-        )
-        return output_text
+        # Convert inference output to JSON
+        return json.dumps(inference_output)
 
     def __call__(self, request):
         try:
             # Parse the incoming request data
             request_data = json.loads(request)
-            
+
             # Preprocess the input data
             inputs = self.preprocess(request_data)
-            
-            # Perform inference
-            outputs = self.inference(inputs)
-            
-            # Postprocess the outputs and return results
-            result = self.postprocess(outputs)
-            return json.dumps({"result": result})
+
+            # Perform inference using the VQA pipeline
+            result = self.inference(inputs)
+
+            # Postprocess the result and return JSON output
+            return self.postprocess(result)
         except Exception as e:
             return json.dumps({"error": str(e)})