3216

handler.py

@@ -47,15 +47,24 @@ class EndpointHandler:
         pixel_values = None
 
         if image_b64:
-            # 关键改动：base64解码
             image_bytes = base64.b64decode(image_b64)
-            pixel_values = self._load_image(image_bytes).unsqueeze(0).to(self.device)
+            pixel_values = self._load_image(image_bytes).unsqueeze(0)
+
+            # ⭐️ 如果模型是 float16，就把输入也变成 half
+            if self.device == "cuda":
+                pixel_values = pixel_values.half()
+
+            pixel_values = pixel_values.to(self.device)
             prompt = prompt.replace("<image>", DEFAULT_IMG_START_TOKEN + DEFAULT_IMG_END_TOKEN)
 
         elif video_b64:
-            # 关键改动：base64解码
             video_bytes = base64.b64decode(video_b64)
-            pixel_values = self._load_video(video_bytes).unsqueeze(0).to(self.device)
+            pixel_values = self._load_video(video_bytes).unsqueeze(0)
+
+            if self.device == "cuda":
+                pixel_values = pixel_values.half()
+
+            pixel_values = pixel_values.to(self.device)
             prompt = prompt.replace("<video>", DEFAULT_VIDEO_START_TOKEN + DEFAULT_VIDEO_END_TOKEN)
 
         return {