model_utils.py

"""
万物有灵 - 模型加载与推理工具
使用 MiniCPM-o 4.5 实现物体识别和人格化对话
"""

import torch
from PIL import Image
from transformers import AutoModel, AutoTokenizer
from personalities import (
    build_identify_prompt,
    build_personality_system_prompt,
    build_followup_system_prompt,
    parse_identification,
    PERSONALITY_TYPES,
    OBJECT_PERSONALITY_HINTS,
    DEFAULT_PERSONALITY,
)
import re

# 全局模型引用
_model = None
_tokenizer = None
_device = None


def get_device():
    """获取可用设备"""
    if torch.cuda.is_available():
        return "cuda"
    return "cpu"


def load_model(model_path: str = "openbmb/MiniCPM-o-4_5", enable_tts: bool = False):
    """
    加载 MiniCPM-o 4.5 模型

    Args:
        model_path: 模型路径或 HuggingFace model ID
        enable_tts: 是否启用 TTS 语音合成（需要更多显存）
    """
    global _model, _tokenizer, _device

    _device = get_device()

    print(f"[Model] Loading model from {model_path} on {_device}...")

    _tokenizer = AutoTokenizer.from_pretrained(
        model_path, trust_remote_code=True
    )

    _model = AutoModel.from_pretrained(
        model_path,
        trust_remote_code=True,
        attn_implementation="sdpa",
        torch_dtype=torch.bfloat16,
        init_audio=enable_tts,
        init_tts=enable_tts,
    )

    _model = _model.to(device=_device, dtype=torch.bfloat16)
    _model.eval()

    if not enable_tts:
        print("[Model] Audio/TTS modules not loaded (save VRAM)")
        torch.cuda.empty_cache()

    print(f"[Model] Model loaded successfully on {_device}")
    return _model, _tokenizer


def is_model_loaded() -> bool:
    """检查模型是否已加载"""
    return _model is not None and _tokenizer is not None


def identify_object(image: Image.Image) -> dict:
    """
    识别图片中的物体

    Args:
        image: PIL Image 对象

    Returns:
        包含物体信息的字典
    """
    if not is_model_loaded():
        return {"error": "模型未加载，请稍候..."}

    prompt = build_identify_prompt()
    msgs = [{"role": "user", "content": [image, prompt]}]

    try:
        with torch.no_grad():
            response = _model.chat(
                image=None,
                msgs=msgs,
                tokenizer=_tokenizer,
                sampling=True,
                temperature=0.5,
                max_new_tokens=512,
            )

        result = parse_identification(response)
        result["raw_response"] = response
        return result

    except Exception as e:
        print(f"[Model] Error during identification: {e}")
        return {"error": str(e), "object_name": "未知物体"}


def chat_with_object(
    image: Image.Image,
    object_info: dict,
    personality_type: str,
    chat_history: list,
    user_message: str,
    is_first_message: bool = False,
) -> str:
    """
    与物体对话

    Args:
        image: 物体的图片
        object_info: 物体识别信息
        personality_type: 性格类型
        chat_history: 之前的对话历史 [{"role": "user/assistant", "content": "..."}]
        user_message: 用户当前消息

    Returns:
        模型的回复
    """
    if not is_model_loaded():
        return "模型未加载，请稍候..."

    object_name = object_info.get("object_name", "未知物体")
    appearance = object_info.get("appearance", "")
    scene = object_info.get("scene", "")
    suggestion = object_info.get("suggestion", "")

    # 构建 system prompt
    if is_first_message or len(chat_history) == 0:
        # 第一次对话：让物体自我介绍
        system_prompt = build_personality_system_prompt(
            object_name=object_name,
            object_appearance=appearance,
            object_scene=scene,
            personality_type=personality_type,
            personality_suggestion=suggestion,
        )
    else:
        # 后续对话：简洁版 system prompt
        id_info = f"外观：{appearance}，场景：{scene}"
        system_prompt = build_followup_system_prompt(
            object_name=object_name,
            personality_type=personality_type,
            identification_info=id_info,
        )

    # 构建消息列表
    msgs = [{"role": "system", "content": system_prompt}]

    # 添加历史对话（第一条用户消息附带图片）
    for i, msg in enumerate(chat_history):
        if msg["role"] == "user":
            if i == 0:
                msgs.append({"role": "user", "content": [image, msg["content"]]})
            else:
                msgs.append({"role": "user", "content": msg["content"]})
        else:
            msgs.append({"role": "assistant", "content": msg["content"]})

    # 添加当前用户消息
    if is_first_message or len(chat_history) == 0:
        # 首次对话：附带图片的自我介绍请求
        msgs.append({"role": "user", "content": [image, "你好！你是谁？请用你的方式自我介绍一下。"]})
    else:
        msgs.append({"role": "user", "content": user_message})

    try:
        with torch.no_grad():
            response = _model.chat(
                image=None,
                msgs=msgs,
                tokenizer=_tokenizer,
                sampling=True,
                temperature=0.8,
                top_p=0.9,
                max_new_tokens=512,
            )

        # 清理响应中的特殊标记
        response = _clean_response(response)
        return response

    except Exception as e:
        print(f"[Model] Error during chat: {e}")
        return f"呜...我好像卡住了（{e}）"


def _clean_response(text: str) -> str:
    """清理模型响应中的特殊标记"""
    # 移除 <box> 标记
    text = re.sub(r"<box>.*?</box>", "", text)
    text = text.replace("<ref>", "").replace("</ref>", "")
    text = text.replace("<box>", "").replace("</box>", "")
    # 移除思考过程标记
    text = re.sub(r"<think.*?>.*?</think.*?>", "", text, flags=re.DOTALL)
    # 清理多余空白
    text = text.strip()
    return text


def suggest_personality(object_name: str) -> str:
    """
    根据物体名称建议性格类型

    Args:
        object_name: 物体名称

    Returns:
        建议的性格类型
    """
    for key, personality in OBJECT_PERSONALITY_HINTS.items():
        if key in object_name:
            return personality
    return DEFAULT_PERSONALITY


def get_model_info() -> str:
    """获取模型状态信息"""
    if not is_model_loaded():
        return "模型未加载"

    device_name = "CPU"
    if _device == "cuda":
        device_name = torch.cuda.get_device_name(0)
        vram = torch.cuda.get_device_properties(0).total_memory / 1024**3
        device_name = f"{device_name} ({vram:.1f}GB)"

    return f"MiniCPM-o 4.5 | 设备: {device_name}"