Add custom handler and requirements

handler.py

@@ -0,0 +1,73 @@
+from typing import Dict, Any
+import base64
+import tempfile
+import os
+import sys
+
+# 确保能导入 videollama2 模块（模型代码需要放同目录或已安装）
+sys.path.append('./')
+
+from videollama2 import model_init, mm_infer
+from videollama2.utils import disable_torch_init
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        # 关闭torch自动初始化，避免重复加载
+        disable_torch_init()
+        # 模型路径，如果HF环境传入的path为空，就用默认的官方仓库地址
+        self.model_path = path or "DAMO-NLP-SG/VideoLLaMA2-7B-16F"
+        # 加载模型、处理器、分词器
+        self.model, self.processor, self.tokenizer = model_init(self.model_path)
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        期待输入数据格式：
+        {
+            "video": "<base64字符串>",  # 视频文件base64编码
+            "prompt": "描述视频内容的自然语言指令"
+        }
+        或者
+        {
+            "image": "<base64字符串>",  # 图片文件base64编码
+            "prompt": "描述图片内容的自然语言指令"
+        }
+        """
+        # 判断输入模态
+        if "video" in data:
+            modal = "video"
+            file_b64 = data["video"]
+        elif "image" in data:
+            modal = "image"
+            file_b64 = data["image"]
+        else:
+            return {"error": "请求必须包含 'video' 或 'image' 字段"}
+
+        prompt = data.get("prompt", "Describe the content.")
+
+        # 临时写入二进制文件，供 processor 读取
+        suffix = ".mp4" if modal == "video" else ".png"
+        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:
+            tmp_file.write(base64.b64decode(file_b64))
+            tmp_path = tmp_file.name
+
+        try:
+            # 处理输入，调用模型推理
+            inputs = self.processor[modal](tmp_path)
+            output = mm_infer(
+                inputs,
+                prompt,
+                model=self.model,
+                tokenizer=self.tokenizer,
+                do_sample=False,
+                modal=modal
+            )
+        finally:
+            # 清理临时文件
+            os.remove(tmp_path)
+
+        # 返回结构统一，方便调用方解析
+        return {
+            "modal": modal,
+            "prompt": prompt,
+            "result": output
+        }

requirements.txt

@@ -0,0 +1,39 @@
+--extra-index-url https://download.pytorch.org/whl/cu118
+# basic dependencies
+torch==2.2.0
+torchvision==0.17.0
+transformers==4.40.0
+tokenizers==0.19.1
+deepspeed==0.13.1
+accelerate==0.26.1
+peft==0.4.0
+timm==1.0.3
+numpy==1.24.4
+# data processing
+decord==0.6.0
+imageio==2.34.0
+imageio-ffmpeg==0.4.9
+moviepy==1.0.3
+opencv-python==4.6.0.66
+pysubs2
+# misc
+scikit-learn==1.2.2
+huggingface_hub==0.23.4
+sentencepiece==0.1.99
+shortuuid
+einops==0.6.1
+einops-exts==0.0.4
+bitsandbytes==0.43.0
+pydantic>=2.0
+markdown2[all]
+gradio==3.50.0
+gradio_client==0.6.1
+httpx==0.24.1
+requests
+openai
+uvicorn
+fastapi
+tensorboard
+wandb
+tabulate
+spaces==0.29.2