Update handler.py

handler.py

@@ -4,15 +4,18 @@ from PIL import Image
 import io
 import base64
 import requests
+import torch
 
 class EndpointHandler():
     def __init__(self, path=""):
         self.processor = AutoProcessor.from_pretrained(path)
-        self.model = Qwen2VLForConditionalGeneration.from_pretrained(path)
+        self.model = Qwen2VLForConditionalGeneration.from_pretrained(
+            path, device_map="auto"
+        )
 
     def __call__(self, data: Any) -> Dict[str, Any]:
         image_input = data.get('image')
-        text_input = data.get('text', "Describe this image.") 
+        text_input = data.get('text', "Describe this image.")
 
         if image_input is None:
             return {"error": "No image provided."}
@@ -26,27 +29,37 @@ class EndpointHandler():
         except Exception as e:
             return {"error": f"Failed to process the image. Details: {str(e)}"}
 
-        messages = [
+        conversation = [
             {
                 "role": "user",
                 "content": [
-                    {"type": "image", "image": image},
+                    {"type": "image"},
                     {"type": "text", "text": text_input},
                 ],
             }
         ]
 
-        text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        text_prompt = self.processor.apply_chat_template(
+            conversation, add_generation_prompt=True
+        )
+
         inputs = self.processor(
-            text=[text],
+            text=[text_prompt],
             images=[image],
             padding=True,
             return_tensors="pt",
         )
 
-        generate_ids = self.model.generate(inputs.input_ids, max_length=30)
+        inputs = inputs.to(self.model.device)
+
+        output_ids = self.model.generate(**inputs, max_new_tokens=128)
+
+        generated_ids = [
+            output_id[len(input_id):] for input_id, output_id in zip(inputs.input_ids, output_ids)
+        ]
+
         output_text = self.processor.batch_decode(
-            generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
+            generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
         )[0]
 
         return {"generated_text": output_text}