test

handler.py

@@ -1,3 +1,4 @@
+### ✅ handler.py（优化版）
 import os
 import torch
 import base64
@@ -28,41 +29,40 @@ class EndpointHandler:
             bos_token_id=1,
             do_sample=False,
             temperature=0.7,
-            max_new_tokens=10240
+            max_new_tokens=1024
         )
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
-        inputs = self.preprocess(data)
-        prediction = self.inference(inputs)
-        return self.postprocess(prediction)
+        try:
+            inputs = self.preprocess(data)
+            prediction = self.inference(inputs)
+            return self.postprocess(prediction)
+        except Exception as e:
+            return {"output": f"❌ 推理失败: {str(e)}"}
 
     def preprocess(self, request: Dict[str, Any]) -> Dict[str, Any]:
         prompt = request["inputs"]
         image_b64 = request.get("image", None)
         video_b64 = request.get("video", None)
+        num_segments = request.get("num_segments", 16)
 
         pixel_values = None
 
         if image_b64:
             image_bytes = base64.b64decode(image_b64)
-            pixel_values = self._load_image(image_bytes).unsqueeze(0)  # [1, 3, 224, 224]
-            if self.device == "cuda":
-                pixel_values = pixel_values.half()
+            pixel_values = self._load_image(image_bytes).unsqueeze(0)
+            pixel_values = pixel_values.half() if self.device == "cuda" else pixel_values
             pixel_values = pixel_values.to(self.device)
             prompt = prompt.replace("<image>", DEFAULT_IMG_START_TOKEN + DEFAULT_IMG_END_TOKEN)
 
         elif video_b64:
             video_bytes = base64.b64decode(video_b64)
-            pixel_values = self._load_video(video_bytes)
-            if self.device == "cuda":
-                pixel_values = pixel_values.half()
+            pixel_values = self._load_video(video_bytes, num_segments)
+            pixel_values = pixel_values.half() if self.device == "cuda" else pixel_values
             pixel_values = pixel_values.to(self.device)
             prompt = prompt.replace("<video>", DEFAULT_VIDEO_START_TOKEN + DEFAULT_VIDEO_END_TOKEN)
 
-        return {
-            "prompt": prompt,
-            "pixel_values": pixel_values
-        }
+        return {"prompt": prompt, "pixel_values": pixel_values}
 
     def inference(self, inputs: Dict[str, Any]) -> str:
         prompt = inputs["prompt"]
@@ -72,37 +72,28 @@ class EndpointHandler:
         model_inputs.pop("token_type_ids", None)
         model_inputs = {k: v.to(self.device) for k, v in model_inputs.items()}
 
+        print("📌 prompt token长度:", model_inputs["input_ids"].shape[1])
         if pixel_values is not None:
-            output = self.model.generate(
-                **model_inputs,
-                pixel_values=pixel_values,
-                generation_config=self.gen_config,
-                return_dict_in_generate=True,
-                output_scores=True
-            )
-        else:
-            output = self.model.language_model.generate(
-                **model_inputs,
-                generation_config=self.gen_config,
-                return_dict_in_generate=True,
-                output_scores=True
-            )
-        # 🧠 打印 debug 信息
+            print("🎞️ pixel shape:", pixel_values.shape)
+
+        output = self.model.generate(
+            **model_inputs,
+            pixel_values=pixel_values,
+            generation_config=self.gen_config,
+            return_dict_in_generate=True,
+            output_scores=True
+        )
+
         generated_ids = output.sequences[0]
-        print("📍生成的 token ids：", generated_ids.tolist())
-        raw_text = self.tokenizer.decode(generated_ids, skip_special_tokens=False)
         clean_text = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
-        print("🧾 带特殊符号的输出：", raw_text)
-        print("✅ 去掉特殊符号的输出：", clean_text)
+        return clean_text
 
-        return clean_text  # 返回干净版本
     def postprocess(self, output: str) -> Dict[str, str]:
         return {"output": output.strip()}
 
     def _load_image(self, image_bytes: bytes) -> torch.Tensor:
         image = Image.open(BytesIO(image_bytes)).convert('RGB')
-        crop_pct = 224 / 256
-        size = int(224 / crop_pct)
+        size = int(224 / (224 / 256))
         transform = T.Compose([
             T.Resize(size, interpolation=InterpolationMode.BICUBIC),
             T.CenterCrop(224),
@@ -111,7 +102,7 @@ class EndpointHandler:
         ])
         return transform(image)
 
-    def _load_video(self, video_bytes: bytes, num_segments=8) -> torch.Tensor:
+    def _load_video(self, video_bytes: bytes, num_segments=16) -> torch.Tensor:
         with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
             tmpfile.write(video_bytes)
             video_path = tmpfile.name
@@ -125,12 +116,10 @@ class EndpointHandler:
             T.Resize(224, interpolation=InterpolationMode.BICUBIC),
             T.CenterCrop(224),
             T.ToTensor(),
-            T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+            T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
         ])
-        processed = [transform(frame) for frame in frames]  # each: [3, 224, 224]
-        video_tensor = torch.stack(processed, dim=0)  # [T, 3, 224, 224]
-        video_tensor = video_tensor.permute(1, 0, 2, 3)  # [3, T, 224, 224]
-        video_tensor = video_tensor.unsqueeze(0)  # [1, 3, T, 224, 224] ✅
+        processed = [transform(frame) for frame in frames]
+        video_tensor = torch.stack(processed, dim=0).permute(1, 0, 2, 3).unsqueeze(0)
         return video_tensor
 
     def get_index(self, num_frames: int, num_segments: int):