src/character_spike/schema.py

from __future__ import annotations

import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any


EXPRESSIONS = ["idle", "listening", "thinking", "worried", "smile", "happy", "talk", "focus"]
CANONICAL_STAGE_SIZE = (900, 1200)
WARM_FOUR_IMAGE_LIMIT_SECONDS = 60.0
EIGHT_ASSET_LIMIT_SECONDS = 180.0
MIN_USABLE_ASSET_COUNT = 6


@dataclass(frozen=True)
class ModelCandidate:
    id: str
    label: str
    family: str
    model_id: str
    mode: str
    default_steps: int
    default_gpu: str
    implemented: bool
    notes: str


MODEL_CANDIDATES: tuple[ModelCandidate, ...] = (
    ModelCandidate(
        id="flux_schnell",
        label="FLUX.1-schnell",
        family="flux",
        model_id="black-forest-labs/FLUX.1-schnell",
        mode="text_to_image_speed_baseline",
        default_steps=4,
        default_gpu="H100",
        implemented=True,
        notes="Speed baseline for main visual candidates.",
    ),
    ModelCandidate(
        id="qwen_image",
        label="Qwen-Image",
        family="qwen_image",
        model_id="Qwen/Qwen-Image",
        mode="text_to_image_quality_candidate",
        default_steps=50,
        default_gpu="H100",
        implemented=True,
        notes="Quality and Chinese prompt candidate; likely slower than FLUX.",
    ),
    ModelCandidate(
        id="qwen_image_edit",
        label="Qwen-Image-Edit",
        family="qwen_image_edit",
        model_id="Qwen/Qwen-Image-Edit",
        mode="instruction_image_edit",
        default_steps=50,
        default_gpu="H100",
        implemented=True,
        notes="Expression and local edit candidate based on a reference image.",
    ),
    ModelCandidate(
        id="qwen_controlnet_union",
        label="Qwen-Image-ControlNet-Union",
        family="qwen_controlnet",
        model_id="InstantX/Qwen-Image-ControlNet-Union",
        mode="pose_canny_depth_control",
        default_steps=30,
        default_gpu="H100",
        implemented=True,
        notes="Structure control candidate for action poses.",
    ),
    ModelCandidate(
        id="instantid_sdxl",
        label="InstantID SDXL",
        family="instantid",
        model_id="InstantX/InstantID",
        mode="identity_preserving_candidate",
        default_steps=30,
        default_gpu="H100",
        implemented=False,
        notes="Tracked as identity-preserving candidate; remote runner is intentionally not enabled until antelopev2/model download path is decided.",
    ),
)


def candidate_by_id(candidate_id: str) -> ModelCandidate:
    for candidate in MODEL_CANDIDATES:
        if candidate.id == candidate_id:
            return candidate
    known = ", ".join(candidate.id for candidate in MODEL_CANDIDATES)
    raise ValueError(f"unknown model candidate: {candidate_id}; expected one of {known}")


def slugify_identifier(value: str, fallback: str = "character") -> str:
    normalized = value.strip().lower()
    normalized = re.sub(r"[^a-z0-9_\-\u4e00-\u9fff]+", "_", normalized)
    normalized = re.sub(r"_+", "_", normalized).strip("_-")
    return normalized or fallback


def project_root() -> Path:
    return Path(__file__).resolve().parents[2]


def default_character_package(character_id: str, display_name: str) -> dict[str, Any]:
    return {
        "id": character_id,
        "name": display_name,
        "display_name": display_name,
        "summary": "自动化角色生成风险验证用原创角色草案。",
        "description": f"{display_name} 是用于验证多表情虚拟角色生成流水线的原创角色。",
        "personality": "冷静、温柔、有清晰边界。",
        "scenario": "用户正在通过虚拟角色实验台与角色进行对话和视觉资产测试。",
        "first_mes": "我在。现在可以开始验证角色生成流程。",
        "alternate_greetings": ["测试频道已接入。", "角色资产验证准备完成。"],
        "mes_example": "",
        "creator_notes": "由自动化角色生成 spike 创建；用于技术验证，不代表最终角色设定。",
        "tags": ["生成测试", "原创角色", "技术验证"],
        "profile": {
            "identity": "自动化角色生成风险验证用原创虚拟角色",
            "core_traits": ["冷静", "温柔", "边界清晰"],
            "relationship_to_user": "把用户当成共同验证系统的协作者",
            "boundaries": ["不声称自己是商业 IP 角色", "不复述商业 IP 官方设定"],
        },
        "dialogue_style": {
            "tone": "自然、简短、清晰",
            "sentence_shape": "中短句",
            "catchphrases": ["我在。"],
        },
        "skills": ["daily_chat", "style_guard"],
        "voice": {"voice_id": "default", "pace": "normal", "energy": 0.5},
        "visual": {
            "accent": "#67e8f9",
            "background": "#111827",
            "background_image": f"{character_id}_spike_background",
            "avatar": character_id,
        },
        "metadata": {"source": "character_generation_spike"},
    }