aaaaaaaa

handler.py

@@ -6,13 +6,10 @@ from io import BytesIO
 from typing import Dict, Any
 from transformers import LlamaTokenizer, GenerationConfig
 from robohusky.model.modeling_husky_embody2 import HuskyForConditionalGeneration
-from robohusky.video_transformers import (
-    GroupNormalize, GroupScale, GroupCenterCrop,
-    Stack, ToTorchFormatTensor, get_index
-)
 from decord import VideoReader, cpu
 import torchvision.transforms as T
 from torchvision.transforms.functional import InterpolationMode
+import tempfile
 
 DEFAULT_IMG_START_TOKEN = "<img>"
 DEFAULT_IMG_END_TOKEN = "</img>"
@@ -48,22 +45,17 @@ class EndpointHandler:
 
         if image_b64:
             image_bytes = base64.b64decode(image_b64)
-            pixel_values = self._load_image(image_bytes).unsqueeze(0)
-
-            # ⭐️ 如果模型是 float16，就把输入也变成 half
+            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).unsqueeze(0)
-
+            pixel_values = self._load_video(video_bytes).unsqueeze(0)  # [1, T, 3, 224, 224]
             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)
 
@@ -114,17 +106,27 @@ class EndpointHandler:
         return transform(image)
 
     def _load_video(self, video_bytes: bytes, num_segments=8) -> torch.Tensor:
-        with open("/tmp/temp_video.mp4", "wb") as f:
-            f.write(video_bytes)
-        vr = VideoReader("/tmp/temp_video.mp4", ctx=cpu(0))
-        frame_indices = get_index(len(vr), num_segments)
-        frames = [Image.fromarray(vr[idx].asnumpy()) for idx in frame_indices]
+        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
+            tmpfile.write(video_bytes)
+            video_path = tmpfile.name
+
+        vr = VideoReader(video_path, ctx=cpu(0))
+        total_frames = len(vr)
+        indices = self.get_index(total_frames, num_segments)
+        frames = [Image.fromarray(vr[i].asnumpy()) for i in indices]
 
         transform = T.Compose([
-            GroupScale(224),
-            GroupCenterCrop(224),
-            Stack(),
-            ToTorchFormatTensor(),
-            GroupNormalize([0.48145466, 0.4578275, 0.40821073], [0.26862954, 0.26130258, 0.27577711])
+            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)),
         ])
-        return transform(frames)
+        processed = [transform(frame) for frame in frames]  # each is [3, 224, 224]
+        video_tensor = torch.stack(processed, dim=0)  # [T, 3, 224, 224]
+        return video_tensor
+
+    def get_index(self, num_frames: int, num_segments: int):
+        if num_frames < num_segments:
+            return list(range(num_frames)) + [num_frames - 1] * (num_segments - num_frames)
+        interval = num_frames / num_segments
+        return [int(interval * i) for i in range(num_segments)]