555

handler.py

@@ -1,4 +1,3 @@
-### ✅ handler.py（优化版）
 import os
 import torch
 import base64
@@ -29,40 +28,41 @@ class EndpointHandler:
             bos_token_id=1,
             do_sample=False,
             # temperature=0.7,
-            max_new_tokens=1024
+            max_new_tokens=4096
         )
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
-        try:
-            inputs = self.preprocess(data)
-            prediction = self.inference(inputs)
-            return self.postprocess(prediction)
-        except Exception as e:
-            return {"output": f"❌ 推理失败: {str(e)}"}
+        inputs = self.preprocess(data)
+        prediction = self.inference(inputs)
+        return self.postprocess(prediction)
 
     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)
-            pixel_values = pixel_values.half() if self.device == "cuda" else pixel_values
+            pixel_values = self._load_image(image_bytes).unsqueeze(0)  # [1, 3, 224, 224]
+            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:
             video_bytes = base64.b64decode(video_b64)
-            pixel_values = self._load_video(video_bytes, num_segments)
-            pixel_values = pixel_values.half() if self.device == "cuda" else pixel_values
+            pixel_values = self._load_video(video_bytes)
+            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 {"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,28 +72,38 @@ 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:
-            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
-        )
-
+            output = self.model.generate(
+                **model_inputs,
+                pixel_values=pixel_values,
+                max_new_tokens=self.gen_config.max_new_tokens,  # 👈 显式传入
+                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 信息
         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)
-        return clean_text
+        print("🧾 带特殊符号的输出：", raw_text)
+        print("✅ 去掉特殊符号的输出：", 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')
-        size = int(224 / (224 / 256))
+        crop_pct = 224 / 256
+        size = int(224 / crop_pct)
         transform = T.Compose([
             T.Resize(size, interpolation=InterpolationMode.BICUBIC),
             T.CenterCrop(224),
@@ -116,10 +126,12 @@ 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]
-        video_tensor = torch.stack(processed, dim=0).permute(1, 0, 2, 3).unsqueeze(0)
+        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] ✅
         return video_tensor
 
     def get_index(self, num_frames: int, num_segments: int):