如何对部署在SageMaker的Llava Llama多模态模型执行推理？

问题概述 #

已通过HuggingFace部署卡片将LLaVA-Llama-2多模态模型部署到SageMaker端点，但部署时指定HF_TASK为text-generation，导致无法处理图文混合输入。调用端点时触发ParamValidationError，提示请求Body必须为字节类型而非字典；核心需求是实现同时传递图片（需按336x336尺寸预处理）和文本提示的推理逻辑。

错误原因 #

1. 直接使用Predictor调用时未对字典格式的请求体做序列化，SageMaker端点要求请求体为字节类型。
2. 默认HuggingFace推理容器仅支持纯文本的text-generation任务，无法处理LLaVA所需的图文输入解析、图片预处理及多模态提示组装逻辑。

解决方案 #

步骤1：编写自定义推理脚本（inference.py） #

实现图文输入解析、图片预处理、提示模板组装及多模态推理逻辑：

import base64
import io
import torch
import json
from PIL import Image
from llava.constants import IMAGE_TOKEN_INDEX
from llava.conversation import conv_templates
from llava.model.builder import load_pretrained_model
from llava.utils import disable_torch_init

# 初始化模型与处理器
disable_torch_init()
model_name = "liuhaotian/llava-llama-2-13b-chat-lightning-preview"
model, vis_processor, text_processor = load_pretrained_model(
    model_name, "cpu", "cuda" if torch.cuda.is_available() else "cpu"
)

def input_fn(request_body, request_content_type):
    # 解析JSON请求，提取Base64图片和文本提示
    input_data = json.loads(request_body)
    # 解码Base64图片并转为PIL格式
    image_bytes = base64.b64decode(input_data["image"])
    image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
    # 按模型要求预处理图片（自动缩放、裁剪至336x336）
    image_tensor = vis_processor(image).unsqueeze(0)
    # 组装LLaVA专用提示模板
    conv = conv_templates["llava_v1"].copy()
    conv.append_message(conv.roles[0], f"{text_processor(input_data['prompt'])}")
    conv.append_message(conv.roles[1], None)
    prompt = conv.get_prompt()
    return {"image": image_tensor, "prompt": prompt}

def predict_fn(input_data, model):
    image_tensor = input_data["image"]
    prompt = input_data["prompt"]
    # 构造带图片标记的输入ID
    input_ids = text_processor([prompt])
    input_ids = torch.cat([
        input_ids[:, :-1], 
        torch.tensor([[IMAGE_TOKEN_INDEX]]), 
        input_ids[:, -1:]
    ], dim=1).to(model.device)
    attention_mask = torch.ones_like(input_ids)
    image_sizes = [(image_tensor.shape[2], image_tensor.shape[3])]
    # 执行推理
    with torch.no_grad():
        output_ids = model.generate(
            images=image_tensor,
            input_ids=input_ids,
            attention_mask=attention_mask,
            image_sizes=image_sizes,
            max_new_tokens=512,
            temperature=0.2
        )
    # 解码输出文本
    output_text = text_processor.decode(output_ids[0], skip_special_tokens=True).strip()
    return {"generated_text": output_text}

def output_fn(prediction, response_content_type):
    return json.dumps(prediction), response_content_type

步骤2：创建依赖文件（requirements.txt） #

添加LLaVA及相关依赖，确保容器能正常运行推理脚本：

llava @ git+https://github.com/haotian-liu/LLaVA.git
pillow
torch>=1.13.1
transformers>=4.26.0
accelerate

步骤3：重新部署模型到SageMaker #

将inference.py和requirements.txt打包为model-artifact.tar.gz，上传至你的S3存储桶，然后修改部署代码：

import sagemaker
import boto3
from sagemaker.huggingface import HuggingFaceModel

try:
    role = sagemaker.get_execution_role()
except ValueError:
    iam = boto3.client('iam')
    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']

# 替换为你上传的模型压缩包的S3路径
model_data = "s3://your-bucket/path/to/model-artifact.tar.gz"

huggingface_model = HuggingFaceModel(
    model_data=model_data,
    transformers_version='4.26.0',
    pytorch_version='1.13.1',
    py_version='py39',
    role=role,
    env={
        'HF_MODEL_ID': 'liuhaotian/llava-llama-2-13b-chat-lightning-preview',
        'HF_TASK': 'multimodal-generation'
    }
)

# 建议使用带GPU的实例，13B模型需至少16GB显存
predictor = huggingface_model.deploy(
    initial_instance_count=1,
    instance_type='ml.g4dn.xlarge'
)

步骤4：正确调用端点推理 #

修改推理代码，将请求序列化为字节格式，同时传递Base64图片和文本提示：

from sagemaker.predictor import Predictor
from base64 import b64encode
import json

endpoint = 'your-endpoint-name'
ENCODING = "utf-8"
IMAGE_NAME = "eiffel_tower_336.jpg"

# 读取图片并转为Base64字符串
with open(IMAGE_NAME, 'rb') as f:
    base64_string = b64encode(f.read()).decode(ENCODING)

# 构造请求体
payload = {
    "prompt": "详细描述这张图片的内容",
    "image": base64_string
}
# 序列化为字节格式
payload_bytes = json.dumps(payload).encode(ENCODING)

predictor = Predictor(endpoint)
inference_response = predictor.predict(data=payload_bytes)
# 解析并打印结果
response_dict = json.loads(inference_response)
print(response_dict["generated_text"])

关键细节说明 #

• 图片预处理：LLaVA的vis_processor会自动完成缩放、裁剪、填充至336x336的操作，无需手动处理。
• 特殊token：推理脚本会自动在提示中插入图片标记token，匹配模型训练时的输入格式。
• 实例选择：13B模型需至少16GB显存，推荐使用ml.g4dn.xlarge或更高规格的GPU实例，避免显存不足。

内容的提问来源于stack exchange，提问作者Csaba Toth