178 lines
6.0 KiB
Python
178 lines
6.0 KiB
Python
from __future__ import annotations
|
||
|
||
import os
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
from typing import Optional, List, Tuple
|
||
|
||
|
||
def _get_env_bool(key: str, default: bool) -> bool:
|
||
"""从环境变量获取布尔值"""
|
||
val = os.environ.get(key, "").lower()
|
||
if val in ("1", "true", "yes", "on"):
|
||
return True
|
||
if val in ("0", "false", "no", "off"):
|
||
return False
|
||
return default
|
||
|
||
|
||
def _get_env_float(key: str, default: float) -> float:
|
||
"""从环境变量获取浮点数"""
|
||
val = os.environ.get(key)
|
||
if val:
|
||
try:
|
||
return float(val)
|
||
except ValueError:
|
||
pass
|
||
return default
|
||
|
||
|
||
def _get_env_int(key: str, default: int) -> int:
|
||
"""从环境变量获取整数"""
|
||
val = os.environ.get(key)
|
||
if val:
|
||
try:
|
||
return int(val)
|
||
except ValueError:
|
||
pass
|
||
return default
|
||
|
||
|
||
# 默认模型文件名
|
||
DEFAULT_UNET_NAME = "anima-preview.safetensors"
|
||
DEFAULT_CLIP_NAME = "qwen_3_06b_base.safetensors"
|
||
DEFAULT_VAE_NAME = "qwen_image_vae.safetensors"
|
||
|
||
|
||
@dataclass
|
||
class AnimaToolConfig:
|
||
"""
|
||
Anima 工具配置。
|
||
|
||
所有配置项都支持通过环境变量覆盖：
|
||
- COMFYUI_URL: ComfyUI Web 服务地址（默认 http://127.0.0.1:8188）
|
||
- ANIMATOOL_DOWNLOAD_IMAGES: 是否下载图片到本地（默认 true）
|
||
- ANIMATOOL_OUTPUT_DIR: 图片输出目录
|
||
- ANIMATOOL_TIMEOUT: 生成超时时间（秒，默认 600）
|
||
- ANIMATOOL_POLL_INTERVAL: 轮询间隔（秒，默认 1）
|
||
- ANIMATOOL_TARGET_MP: 目标像素数（MP，默认 1.0）
|
||
- ANIMATOOL_ROUND_TO: 分辨率对齐倍数（默认 16）
|
||
|
||
示例：
|
||
# Windows PowerShell
|
||
$env:COMFYUI_URL = "http://192.168.1.100:8188"
|
||
|
||
# Linux/macOS
|
||
export COMFYUI_URL=http://192.168.1.100:8188
|
||
|
||
# Cursor MCP 配置中
|
||
{
|
||
"mcpServers": {
|
||
"anima-tool": {
|
||
"command": "python",
|
||
"args": ["mcp_server.py"],
|
||
"env": {
|
||
"COMFYUI_URL": "http://192.168.1.100:8188"
|
||
}
|
||
}
|
||
}
|
||
}
|
||
"""
|
||
|
||
# ComfyUI 服务地址
|
||
# 支持：本地 (127.0.0.1)、局域网 (192.168.x.x)、Docker (host.docker.internal)
|
||
comfyui_url: str = field(
|
||
default_factory=lambda: os.environ.get("COMFYUI_URL", "http://127.0.0.1:8188")
|
||
)
|
||
|
||
# 下载模式：把 /view 拿到的图片保存到本地
|
||
download_images: bool = field(
|
||
default_factory=lambda: _get_env_bool("ANIMATOOL_DOWNLOAD_IMAGES", True)
|
||
)
|
||
output_dir: Path = field(
|
||
default_factory=lambda: Path(
|
||
os.environ.get("ANIMATOOL_OUTPUT_DIR", "")
|
||
) if os.environ.get("ANIMATOOL_OUTPUT_DIR") else Path(__file__).resolve().parent.parent / "outputs"
|
||
)
|
||
|
||
# 轮询历史接口等待执行完成
|
||
timeout_s: float = field(
|
||
default_factory=lambda: _get_env_float("ANIMATOOL_TIMEOUT", 600.0)
|
||
)
|
||
poll_interval_s: float = field(
|
||
default_factory=lambda: _get_env_float("ANIMATOOL_POLL_INTERVAL", 1.0)
|
||
)
|
||
|
||
# 分辨率生成：当只给 aspect_ratio 时，按目标像素数估算宽高
|
||
target_megapixels: float = field(
|
||
default_factory=lambda: _get_env_float("ANIMATOOL_TARGET_MP", 1.0)
|
||
)
|
||
# 宽高向上取整到 round_to 的倍数
|
||
# 注意：Anima 基于 Cosmos 架构，VAE 缩放 8 倍后还需被 spatial_patch_size=2 整除
|
||
# 所以必须是 8×2=16 的倍数，否则会报错 "should be divisible by spatial_patch_size"
|
||
round_to: int = field(
|
||
default_factory=lambda: _get_env_int("ANIMATOOL_ROUND_TO", 16)
|
||
)
|
||
|
||
# -------------------------
|
||
# 模型配置
|
||
# -------------------------
|
||
# ComfyUI models 根目录（用于模型预检查）
|
||
# 如果不设置，则尝试自动探测，探测失败则跳过预检查（远程 ComfyUI 场景）
|
||
comfyui_models_dir: Optional[Path] = field(
|
||
default_factory=lambda: (
|
||
Path(os.environ["COMFYUI_MODELS_DIR"]) if os.environ.get("COMFYUI_MODELS_DIR") else
|
||
(Path.cwd() / "models" if (Path.cwd() / "models").exists() else
|
||
(Path(__file__).resolve().parent.parent.parent.parent / "models" if (Path(__file__).resolve().parent.parent.parent.parent / "models").exists() else None))
|
||
)
|
||
)
|
||
|
||
# 模型文件名（可通过环境变量或参数覆盖）
|
||
unet_name: str = field(
|
||
default_factory=lambda: os.environ.get("ANIMATOOL_UNET_NAME", DEFAULT_UNET_NAME)
|
||
)
|
||
clip_name: str = field(
|
||
default_factory=lambda: os.environ.get("ANIMATOOL_CLIP_NAME", DEFAULT_CLIP_NAME)
|
||
)
|
||
vae_name: str = field(
|
||
default_factory=lambda: os.environ.get("ANIMATOOL_VAE_NAME", DEFAULT_VAE_NAME)
|
||
)
|
||
|
||
# 是否启用模型预检查（默认启用，但需要设置 COMFYUI_MODELS_DIR）
|
||
check_models: bool = field(
|
||
default_factory=lambda: _get_env_bool("ANIMATOOL_CHECK_MODELS", True)
|
||
)
|
||
|
||
def get_model_paths(self) -> dict:
|
||
"""
|
||
返回模型文件的预期路径（相对于 ComfyUI models 目录）。
|
||
"""
|
||
return {
|
||
"unet": ("diffusion_models", self.unet_name),
|
||
"clip": ("text_encoders", self.clip_name),
|
||
"vae": ("vae", self.vae_name),
|
||
}
|
||
|
||
def check_models_exist(self) -> Tuple[bool, List[str]]:
|
||
"""
|
||
检查模型文件是否存在。
|
||
|
||
Returns:
|
||
(all_exist, missing_files): 是否全部存在，缺失的文件列表
|
||
"""
|
||
if not self.comfyui_models_dir:
|
||
# 未配置 models 目录，跳过检查
|
||
return True, []
|
||
|
||
models_dir = Path(self.comfyui_models_dir)
|
||
if not models_dir.exists():
|
||
return False, [f"ComfyUI models 目录不存在: {models_dir}"]
|
||
|
||
missing = []
|
||
for model_type, (subdir, filename) in self.get_model_paths().items():
|
||
model_path = models_dir / subdir / filename
|
||
if not model_path.exists():
|
||
missing.append(f"{model_type}: {subdir}/{filename}")
|
||
|
||
return len(missing) == 0, missing
|