app.py

"""
万物有灵 (Animism) - 万物皆有灵，指哪说哪！
用 MiniCPM-o 4.5 赋予任何物体人格和声音

Gradio Web App for HuggingFace Space
"""

import gradio as gr
from PIL import Image
import base64
import io
import os

from personalities import PERSONALITY_TYPES, DEFAULT_PERSONALITY
from model_utils import (
    load_model,
    identify_object,
    chat_with_object,
    get_model_info,
    suggest_personality,
)

# ==================== 配置 ====================

MODEL_PATH = os.environ.get("MODEL_PATH", "openbmb/MiniCPM-o-4_5")
ENABLE_TTS = os.environ.get("ENABLE_TTS", "false").lower() == "true"


# ==================== 应用状态 ====================


class AppState:
    """管理应用状态"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.image = None
        self.object_info = {}
        self.personality_type = DEFAULT_PERSONALITY
        self.chat_history = []
        self.is_identified = False

    def to_dict(self):
        d = {
            "object_info": self.object_info,
            "personality_type": self.personality_type,
            "chat_history": self.chat_history,
            "is_identified": self.is_identified,
        }
        if self.image is not None:
            buf = io.BytesIO()
            self.image.save(buf, format="PNG")
            d["image_b64"] = base64.b64encode(buf.getvalue()).decode("utf-8")
        return d

    def from_dict(self, d):
        if not d:
            self.reset()
            return
        self.object_info = d.get("object_info", {})
        self.personality_type = d.get("personality_type", DEFAULT_PERSONALITY)
        self.chat_history = d.get("chat_history", [])
        self.is_identified = d.get("is_identified", False)
        if d.get("image_b64"):
            self.image = Image.open(io.BytesIO(base64.b64decode(d["image_b64"])))
        else:
            self.image = None


# ==================== 核心逻辑 ====================


_model_init_attempted = False


def init_model():
    """初始化模型（仅尝试一次，失败不阻塞启动）"""
    global _model_init_attempted
    _model_init_attempted = True
    try:
        load_model(model_path=MODEL_PATH, enable_tts=ENABLE_TTS)
        return "ready"
    except Exception as e:
        print(f"[Init] Failed to load model: {e}")
        print("[Init] Will retry when user interacts...")
        return f"error: {e}"


def ensure_model_loaded():
    """确保模型已加载，未加载则尝试加载"""
    from model_utils import is_model_loaded
    if is_model_loaded():
        return True
    try:
        print("[Model] Retrying model load...")
        load_model(model_path=MODEL_PATH, enable_tts=ENABLE_TTS)
        return True
    except Exception as e:
        print(f"[Model] Retry failed: {e}")
        return False


def _parse_personality(radio_value: str) -> str:
    """从 radio 选项值解析性格类型名称"""
    if not radio_value:
        return DEFAULT_PERSONALITY
    for p in PERSONALITY_TYPES:
        if radio_value.startswith(p):
            return p
    return DEFAULT_PERSONALITY


def process_image(image, personality_radio, state_dict, chatbot):
    """处理上传图片：识别物体 + 生成自我介绍"""
    if image is None:
        return chatbot, state_dict, "请先上传或拍摄一张图片", get_model_info()

    # 确保模型已加载（支持启动失败后重试）
    if not ensure_model_loaded():
        chatbot = chatbot + [
            {"role": "assistant", "content": "❌ 模型加载失败，请检查日志或重启应用"}
        ]
        return chatbot, state_dict, "❌ 模型未就绪", get_model_info()

    state = AppState()
    state.from_dict(state_dict)

    # 转换图片
    if isinstance(image, dict):
        image = image.get("composite", image.get("image"))
    if image is None:
        return chatbot, state_dict, "请先上传或拍摄一张图片", get_model_info()
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image).convert("RGB")
    else:
        image = image.convert("RGB")

    state.image = image
    state.personality_type = _parse_personality(personality_radio)
    state.chat_history = []
    state.is_identified = False

    # 识别物体
    object_info = identify_object(image)

    if "error" in object_info:
        error_msg = object_info.get("error", "识别失败")
        chatbot = chatbot + [{"role": "assistant", "content": f"❌ 识别出错了：{error_msg}"}]
        return chatbot, state.to_dict(), f"❌ 识别失败", get_model_info()

    state.object_info = object_info
    state.is_identified = True
    object_name = object_info.get("object_name", "未知物体")

    # 自动建议性格
    auto_personality = suggest_personality(object_name)
    if state.personality_type not in PERSONALITY_TYPES:
        state.personality_type = auto_personality

    persona = PERSONALITY_TYPES[state.personality_type]

    # 信息卡片
    info_text = (
        f"🔍 识别到：**{object_name}**\n\n"
        f"👀 外观：{object_info.get('appearance', '未知')}\n\n"
        f"📍 场景：{object_info.get('scene', '未知')}\n\n"
        f"🎭 性格：{state.personality_type} {persona['emoji']} {persona['description']}"
    )

    # 让物体自我介绍
    intro = chat_with_object(
        image=state.image,
        object_info=state.object_info,
        personality_type=state.personality_type,
        chat_history=[],
        user_message="你好！你是谁？",
        is_first_message=True,
    )

    state.chat_history.append({"role": "user", "content": "你好！你是谁？"})
    state.chat_history.append({"role": "assistant", "content": intro})

    chatbot = chatbot + [
        {"role": "assistant", "content": f"🎭 **{object_name}** 觉醒了！性格：{state.personality_type} {persona['emoji']}\n\n{intro}"}
    ]

    return chatbot, state.to_dict(), info_text, get_model_info()


def send_message(user_message, state_dict, chatbot):
    """发送消息与物体对话"""
    if not user_message or not user_message.strip():
        return chatbot, state_dict, ""

    if not ensure_model_loaded():
        chatbot = chatbot + [
            {"role": "user", "content": user_message},
            {"role": "assistant", "content": "❌ 模型未就绪，请稍后再试"},
        ]
        return chatbot, state_dict, ""

    state = AppState()
    state.from_dict(state_dict)

    if not state.is_identified or state.image is None:
        chatbot = chatbot + [
            {"role": "user", "content": user_message},
            {"role": "assistant", "content": "请先上传图片，让物体觉醒！"},
        ]
        return chatbot, state.to_dict(), ""

    # 调用模型
    response = chat_with_object(
        image=state.image,
        object_info=state.object_info,
        personality_type=state.personality_type,
        chat_history=state.chat_history,
        user_message=user_message,
    )

    # 更新历史
    state.chat_history.append({"role": "user", "content": user_message})
    state.chat_history.append({"role": "assistant", "content": response})

    chatbot = chatbot + [
        {"role": "user", "content": user_message},
        {"role": "assistant", "content": response},
    ]

    return chatbot, state.to_dict(), ""


def switch_personality(personality_radio, state_dict, chatbot):
    """切换性格类型"""
    state = AppState()
    state.from_dict(state_dict)

    if not state.is_identified or state.image is None:
        return chatbot, state.to_dict(), "请先上传图片唤醒物体"

    personality = _parse_personality(personality_radio)
    state.personality_type = personality
    state.chat_history = []

    persona = PERSONALITY_TYPES[personality]
    object_name = state.object_info.get("object_name", "未知物体")

    # 新性格自我介绍
    intro = chat_with_object(
        image=state.image,
        object_info=state.object_info,
        personality_type=personality,
        chat_history=[],
        user_message="你好！你是谁？",
        is_first_message=True,
    )

    state.chat_history.append({"role": "user", "content": "你好！你是谁？"})
    state.chat_history.append({"role": "assistant", "content": intro})

    info_text = (
        f"🔍 识别到：**{object_name}**\n\n"
        f"🎭 性格切换为：{personality} {persona['emoji']} {persona['description']}\n\n"
        f"💬 {intro}"
    )

    chatbot = chatbot + [
        {"role": "assistant", "content": f"🎭 性格切换为：{personality} {persona['emoji']}\n\n{intro}"}
    ]

    return chatbot, state.to_dict(), info_text


def clear_chat(state_dict):
    """清除对话"""
    state = AppState()
    state.from_dict(state_dict)
    state.chat_history = []
    return [], state.to_dict()


def reset_all():
    """完全重置"""
    state = AppState()
    return [], state.to_dict(), "等待上传图片...", None, get_model_info()


# ==================== Gradio UI ====================

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Noto+Sans+SC:wght@300;400;500;700&display=swap');

:root {
    --primary: #6C5CE7;
    --primary-light: #A29BFE;
    --accent: #FD79A8;
}

.gradio-container {
    font-family: 'Noto Sans SC', sans-serif !important;
    max-width: 1200px !important;
}

.app-title {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    font-weight: 700;
    font-size: 2.5rem;
    text-align: center;
    margin: 0;
}

.app-subtitle {
    text-align: center;
    color: #636e72;
    font-size: 1.1rem;
    margin-top: 0.5rem;
}

.object-info-card {
    background: linear-gradient(135deg, #f8f9ff 0%, #f0f4ff 100%) !important;
    border: 1px solid #d1d9ff !important;
    border-radius: 16px !important;
    padding: 20px !important;
    box-shadow: 0 4px 12px rgba(108, 92, 231, 0.1) !important;
}

.quick-btn {
    font-size: 0.85rem !important;
    border-radius: 20px !important;
    margin: 2px !important;
}
"""


def create_app():
    """创建 Gradio 应用"""

    with gr.Blocks(
        title="万物有灵 - Animism",
    ) as app:
        state = gr.State(value=AppState().to_dict())

        # ====== 标题 ======
        gr.HTML(
            """
            <div style="text-align: center; padding: 20px 0;">
                <h1 class="app-title">万物有灵</h1>
                <p class="app-subtitle">✨ 指哪说话，万物成精！对准任何物体，赋予它人格和声音 ✨</p>
                <p style="color: #b2bec3; font-size: 0.85rem;">Powered by MiniCPM-o 4.5 全模态大模型</p>
            </div>
            """
        )

        # ====== 主内容 ======
        with gr.Row(equal_height=True):
            # ---- 左侧：图片 & 设置 ----
            with gr.Column(scale=1, min_width=350):
                gr.Markdown("### 📸 唤醒物体")
                image_input = gr.Image(
                    label="上传或拍摄物体照片",
                    sources=["upload", "webcam"],
                    type="pil",
                    height=300,
                )

                gr.Markdown("### 🎭 选择性格")
                personality_radio = gr.Radio(
                    choices=[
                        f"{p} {PERSONALITY_TYPES[p]['emoji']} {PERSONALITY_TYPES[p]['description']}"
                        for p in PERSONALITY_TYPES
                    ],
                    value=f"傲娇 {PERSONALITY_TYPES['傲娇']['emoji']} {PERSONALITY_TYPES['傲娇']['description']}",
                    label="",
                    interactive=True,
                )

                with gr.Row():
                    awaken_btn = gr.Button("✨ 唤醒！", variant="primary", size="lg")
                    reset_btn = gr.Button("🔄 重置", size="lg")

                object_info = gr.Markdown(
                    "等待上传图片...",
                    elem_classes=["object-info-card"],
                )

                gr.Markdown("### 💡 快捷提问")
                quick_questions = [
                    "你平时都在做什么？",
                    "你有什么烦恼吗？",
                    "你最喜欢什么？",
                    "你觉得人类怎么样？",
                    "讲讲你最难忘的经历",
                ]
                with gr.Row():
                    quick_btn_1 = gr.Button(quick_questions[0], elem_classes=["quick-btn"], size="sm")
                    quick_btn_2 = gr.Button(quick_questions[1], elem_classes=["quick-btn"], size="sm")
                with gr.Row():
                    quick_btn_3 = gr.Button(quick_questions[2], elem_classes=["quick-btn"], size="sm")
                    quick_btn_4 = gr.Button(quick_questions[3], elem_classes=["quick-btn"], size="sm")
                with gr.Row():
                    quick_btn_5 = gr.Button(quick_questions[4], elem_classes=["quick-btn"], size="sm")

            # ---- 右侧：对话区 ----
            with gr.Column(scale=2, min_width=500):
                gr.Markdown("### 💬 与物体对话")
                chatbot = gr.Chatbot(
                    label="",
                    height=520,
                    placeholder="📸 上传图片 → ✨ 点击「唤醒」→ 💬 开始对话",
                )

                with gr.Row():
                    msg_input = gr.Textbox(
                        placeholder="对物体说点什么...",
                        scale=5,
                        show_label=False,
                        lines=1,
                    )
                    send_btn = gr.Button("发送 ➤", variant="primary", scale=1)

                with gr.Row():
                    clear_btn = gr.Button("🗑️ 清空对话", size="sm")
                    switch_btn = gr.Button("🔄 切换性格", size="sm")

                model_status = gr.Markdown(get_model_info())

        # ====== 事件绑定 ======

        # 唤醒
        awaken_btn.click(
            process_image,
            inputs=[image_input, personality_radio, state, chatbot],
            outputs=[chatbot, state, object_info, model_status],
        )

        # 发送消息
        msg_input.submit(
            send_message,
            inputs=[msg_input, state, chatbot],
            outputs=[chatbot, state, msg_input],
        )
        send_btn.click(
            send_message,
            inputs=[msg_input, state, chatbot],
            outputs=[chatbot, state, msg_input],
        )

        # 快捷提问
        for btn, q in zip(
            [quick_btn_1, quick_btn_2, quick_btn_3, quick_btn_4, quick_btn_5],
            quick_questions,
        ):
            btn.click(
                lambda msg, s, cb: send_message(msg, s, cb),
                inputs=[gr.State(q), state, chatbot],
                outputs=[chatbot, state, msg_input],
            )

        # 清空对话
        clear_btn.click(
            clear_chat,
            inputs=[state],
            outputs=[chatbot, state],
        )

        # 切换性格
        switch_btn.click(
            switch_personality,
            inputs=[personality_radio, state, chatbot],
            outputs=[chatbot, state, object_info],
        )

        # 重置
        reset_btn.click(
            reset_all,
            outputs=[chatbot, state, object_info, image_input, model_status],
        )

        # ====== 底部 ======
        gr.HTML(
            """
            <div style="text-align: center; padding: 20px 0; color: #b2bec3; font-size: 0.8rem;">
                <p>万物有灵 (Animism) | Powered by MiniCPM-o 4.5 | Hackathon Project</p>
                <p>🔍 视觉识别 + 🎭 人格赋予 + 💬 对话交互</p>
            </div>
            """
        )

    return app


# ==================== 启动 ====================

if __name__ == "__main__":
    print("=" * 60)
    print("  万物有灵 (Animism) - 万物皆有灵，指哪说哪！")
    print("=" * 60)

    model_status = init_model()
    print(f"[Init] Model status: {model_status}")

    app = create_app()
    app.launch(
        server_name="0.0.0.0",
        server_port=80,
        share=False,
        debug=True,
        css=CSS,
        theme=gr.themes.Soft(
            primary_hue="violet",
            secondary_hue="pink",
            neutral_hue="slate",
        ),
    )