Nakano Kenji
c743ef5981
- README: Features、Parameters 表、HTTP 路由表、目录结构 - wiki/API-Reference: 新增 list_anima_history/reroll_anima_image/list_anima_models 文档 - wiki/Home: Features 列表补充 - wiki/FAQ: 更新批量生成答案，新增 reroll FAQ - mcp_server: generate/reroll 响应追加历史 ID 提示 - anima_executor: generate 返回 history_id 字段
656 lines
25 KiB
Python
656 lines
25 KiB
Python
from __future__ import annotations
|
||
|
||
import base64
|
||
import json
|
||
import math
|
||
import time
|
||
import uuid
|
||
from copy import deepcopy
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional, Tuple
|
||
from urllib.parse import urlencode, urljoin
|
||
|
||
from .config import AnimaToolConfig
|
||
from .history import HistoryManager
|
||
|
||
|
||
def _round_up(x: int, base: int) -> int:
|
||
if base <= 1:
|
||
return x
|
||
return int(math.ceil(x / base) * base)
|
||
|
||
|
||
def _parse_aspect_ratio(ratio: str) -> float:
|
||
"""
|
||
解析 "16:9" -> 16/9。
|
||
"""
|
||
s = (ratio or "").strip()
|
||
if ":" not in s:
|
||
raise ValueError(f"aspect_ratio 必须形如 '16:9'，收到：{ratio!r}")
|
||
a_str, b_str = s.split(":", 1)
|
||
a = float(a_str.strip())
|
||
b = float(b_str.strip())
|
||
if a <= 0 or b <= 0:
|
||
raise ValueError(f"aspect_ratio 两边必须 > 0，收到：{ratio!r}")
|
||
return a / b
|
||
|
||
|
||
def estimate_size_from_ratio(
|
||
*,
|
||
aspect_ratio: str,
|
||
target_megapixels: float = 1.0,
|
||
round_to: int = 16,
|
||
) -> Tuple[int, int]:
|
||
"""
|
||
只给定长宽比时，估算 width/height，使像素数接近 target_megapixels。
|
||
宽高会向上取整到 round_to 的倍数。
|
||
|
||
注意：Anima 基于 Cosmos 架构，VAE 缩放 8 倍后还需被 spatial_patch_size=2 整除，
|
||
所以 round_to 必须至少是 16（8×2），否则会报错 "should be divisible by spatial_patch_size"。
|
||
"""
|
||
r = _parse_aspect_ratio(aspect_ratio)
|
||
target_px = max(1.0, float(target_megapixels)) * 1_000_000.0
|
||
w = int(math.sqrt(target_px * r))
|
||
h = int(math.sqrt(target_px / r))
|
||
w = _round_up(max(64, w), round_to)
|
||
h = _round_up(max(64, h), round_to)
|
||
return w, h
|
||
|
||
|
||
def align_dimension(value: int, round_to: int = 16) -> int:
|
||
"""
|
||
将尺寸对齐到 round_to 的倍数（向上取整）。
|
||
用于处理用户直接传入的 width/height。
|
||
"""
|
||
return _round_up(max(64, int(value)), round_to)
|
||
|
||
|
||
def _join_csv(*parts: str) -> str:
|
||
cleaned: List[str] = []
|
||
for p in parts:
|
||
if p is None:
|
||
continue
|
||
s = str(p).strip()
|
||
if not s:
|
||
continue
|
||
cleaned.append(s)
|
||
return ", ".join(cleaned)
|
||
|
||
|
||
def build_anima_positive_text(prompt_json: Dict[str, Any]) -> str:
|
||
"""
|
||
按 Anima 推荐顺序拼接正面提示词（全部逗号连接，不分行）。
|
||
顺序：[质量/安全/年份] [人数] [角色] [作品] [画师] [风格] [外观] [标签] [环境] [自然语言]
|
||
|
||
说明：nltags（自然语言补充）放在最后，因为它是"实在没法用 tag 才写"的兜底描述。
|
||
"""
|
||
return _join_csv(
|
||
prompt_json.get("quality_meta_year_safe", ""),
|
||
prompt_json.get("count", ""),
|
||
prompt_json.get("character", ""),
|
||
prompt_json.get("series", ""),
|
||
prompt_json.get("artist", ""),
|
||
prompt_json.get("style", ""),
|
||
prompt_json.get("appearance", ""),
|
||
prompt_json.get("tags", ""),
|
||
prompt_json.get("environment", ""),
|
||
prompt_json.get("nltags", ""),
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class GeneratedImage:
|
||
filename: str
|
||
subfolder: str
|
||
folder_type: str
|
||
view_url: str
|
||
saved_path: Optional[str] = None
|
||
content: Optional[bytes] = None # 原始图片数据
|
||
|
||
|
||
class AnimaExecutor:
|
||
_SUPPORTED_MODEL_TYPES = ("loras", "diffusion_models", "vae", "text_encoders")
|
||
"""
|
||
将结构化 JSON 注入 ComfyUI prompt 并执行，获取输出图片。
|
||
"""
|
||
|
||
def __init__(self, config: Optional[AnimaToolConfig] = None):
|
||
self.config = config or AnimaToolConfig()
|
||
self._client_id = str(uuid.uuid4())
|
||
|
||
# 远端 ComfyUI 返回的模型名称分隔符（Windows 常为 "\\"，Linux 常为 "/"）
|
||
self._remote_model_path_sep_cache: Dict[str, str] = {}
|
||
|
||
template_path = Path(__file__).resolve().parent / "workflow_template.json"
|
||
with template_path.open("r", encoding="utf-8") as f:
|
||
self._workflow_template: Dict[str, Any] = json.load(f)
|
||
|
||
# 生成历史管理器
|
||
self.history = HistoryManager()
|
||
|
||
# -------------------------
|
||
# Model listing / metadata
|
||
# -------------------------
|
||
def _read_lora_metadata(self, lora_name: str) -> Optional[Dict[str, Any]]:
|
||
"""读取 LoRA 的 sidecar 元数据文件（同名 .json）。"""
|
||
import os
|
||
if not self.config.comfyui_models_dir:
|
||
return None
|
||
|
||
models_dir = Path(self.config.comfyui_models_dir)
|
||
# 兼容处理：lora_name 可能带路径，针对 Windows 统一斜杠并去除首尾空格
|
||
clean_name = lora_name.strip().replace("/", os.sep).replace("\\", os.sep)
|
||
# 核心修复：如果 clean_name 以斜杠开头，Path / 拼接会变成绝对路径导致失败
|
||
while clean_name.startswith(os.sep):
|
||
clean_name = clean_name[1:]
|
||
|
||
# 尝试几种可能的路径拼接方式
|
||
search_paths = [
|
||
# 1. 标准路径: models/loras/subfolder/file.safetensors.json
|
||
models_dir / "loras" / f"{clean_name}.json",
|
||
# 2. 移除扩展名后的路径: models/loras/subfolder/file.json
|
||
models_dir / "loras" / f"{os.path.splitext(clean_name)[0]}.json",
|
||
# 3. 如果 clean_name 本身已经包含 loras 目录（容错）
|
||
models_dir / f"{clean_name}.json" if "loras" in clean_name.lower() else None
|
||
]
|
||
|
||
for meta_path in search_paths:
|
||
if meta_path and meta_path.exists():
|
||
try:
|
||
return json.loads(meta_path.read_text(encoding="utf-8"))
|
||
except Exception:
|
||
continue
|
||
return None
|
||
|
||
def list_models(self, model_type: str) -> List[Dict[str, Any]]:
|
||
"""列出 ComfyUI 模型文件。
|
||
|
||
- model_type=loras：强制只返回存在 sidecar 元数据（.json）的 LoRA。
|
||
- 其他类型：返回 ComfyUI API 的原始列表。
|
||
- 返回的 name 统一使用正斜杠（/）作为分隔符，避免 Windows 反斜杠转义问题。
|
||
"""
|
||
model_type = (model_type or "").strip()
|
||
if model_type not in self._SUPPORTED_MODEL_TYPES:
|
||
raise ValueError(f"不支持的 model_type={model_type!r}，仅支持：{self._SUPPORTED_MODEL_TYPES}")
|
||
|
||
url = urljoin(self.config.comfyui_url.rstrip("/") + "/", f"models/{model_type}")
|
||
files = self._http_get_json(url)
|
||
|
||
if not isinstance(files, list):
|
||
raise RuntimeError(f"ComfyUI /models/{model_type} 返回异常：{files!r}")
|
||
|
||
results: List[Dict[str, Any]] = []
|
||
for raw_name in files:
|
||
if not isinstance(raw_name, str) or not raw_name.strip():
|
||
continue
|
||
|
||
# 统一使用正斜杠格式返回，避免 Windows 反斜杠在 JSON 中的转义问题
|
||
normalized_name = raw_name.replace("\\", "/")
|
||
item: Dict[str, Any] = {"name": normalized_name}
|
||
|
||
if model_type == "loras":
|
||
# 读取 sidecar 时仍使用原始路径（因为文件系统可能需要系统分隔符）
|
||
meta = self._read_lora_metadata(raw_name)
|
||
if not meta:
|
||
# 强制要求：不提供 json sidecar 的 LoRA 不允许被 list 出来
|
||
continue
|
||
item["metadata"] = meta
|
||
|
||
results.append(item)
|
||
|
||
return results
|
||
|
||
def _detect_remote_model_path_sep(self, model_type: str) -> str:
|
||
"""探测远端 ComfyUI 返回的模型名路径分隔符。
|
||
|
||
ComfyUI 在 Windows 下通常使用 "\\" 返回子目录模型名，在 Linux/macOS 下通常使用 "/"。
|
||
这里通过调用 /models/{model_type} 做一次探测并缓存。
|
||
"""
|
||
model_type = (model_type or "").strip()
|
||
if model_type in self._remote_model_path_sep_cache:
|
||
return self._remote_model_path_sep_cache[model_type]
|
||
|
||
url = urljoin(self.config.comfyui_url.rstrip("/") + "/", f"models/{model_type}")
|
||
files = self._http_get_json(url)
|
||
|
||
sep = "/"
|
||
if isinstance(files, list):
|
||
for it in files:
|
||
if not isinstance(it, str):
|
||
continue
|
||
if "\\" in it:
|
||
sep = "\\"
|
||
break
|
||
if "/" in it:
|
||
sep = "/"
|
||
break
|
||
|
||
self._remote_model_path_sep_cache[model_type] = sep
|
||
return sep
|
||
|
||
def _normalize_remote_model_name(self, name: str, model_type: str) -> str:
|
||
"""将用户输入的 name 规范化为远端 ComfyUI 可接受的模型名格式。"""
|
||
import os
|
||
|
||
s = (name or "").strip()
|
||
if not s:
|
||
return s
|
||
|
||
# 先统一把两种分隔符都视为路径分隔
|
||
remote_sep = self._detect_remote_model_path_sep(model_type)
|
||
s = s.replace("/", remote_sep).replace("\\", remote_sep)
|
||
|
||
# 去掉可能的前导分隔符（否则在某些 Path 拼接/校验场景会变成绝对路径语义）
|
||
while s.startswith(remote_sep):
|
||
s = s[len(remote_sep):]
|
||
# 也顺手把当前系统分隔符的前导去掉
|
||
while s.startswith(os.sep):
|
||
s = s[len(os.sep):]
|
||
|
||
return s
|
||
|
||
# -------------------------
|
||
# Workflow helpers
|
||
# -------------------------
|
||
def _inject_loras(self, wf: Dict[str, Any], loras: Any) -> None:
|
||
"""在 UNET 与 KSampler(model) 之间注入多 LoRA（仅 UNET）。
|
||
|
||
loras: [{"name": "xxx.safetensors", "weight": 0.8}, ...]
|
||
"""
|
||
if not loras:
|
||
return
|
||
if not isinstance(loras, list):
|
||
raise ValueError("loras 必须是数组：[{name, weight}, ...]")
|
||
|
||
# 兼容模板变化：以 KSampler 的 model 输入为起点
|
||
if "19" not in wf or "inputs" not in wf["19"] or "model" not in wf["19"]["inputs"]:
|
||
raise RuntimeError("workflow_template.json 缺少 KSampler(19).inputs.model，无法注入 LoRA")
|
||
|
||
prev_model = wf["19"]["inputs"]["model"]
|
||
|
||
# 生成不会冲突的数字 node id
|
||
numeric_ids = [int(k) for k in wf.keys() if str(k).isdigit()]
|
||
next_id = (max(numeric_ids) + 1) if numeric_ids else 1
|
||
|
||
for i, lora in enumerate(loras):
|
||
if not isinstance(lora, dict):
|
||
continue
|
||
name = str(lora.get("name") or "").strip()
|
||
if not name:
|
||
continue
|
||
# ComfyUI 会对 lora_name 做枚举校验，必须与 /models/loras 返回的字符串完全一致
|
||
name = self._normalize_remote_model_name(name, "loras")
|
||
weight = float(lora.get("weight", 1.0))
|
||
|
||
node_id = str(next_id + i)
|
||
wf[node_id] = {
|
||
"class_type": "LoraLoaderModelOnly",
|
||
"inputs": {
|
||
"model": prev_model,
|
||
"lora_name": name,
|
||
"strength_model": weight,
|
||
},
|
||
}
|
||
prev_model = [node_id, 0]
|
||
|
||
wf["19"]["inputs"]["model"] = prev_model
|
||
|
||
# -------------------------
|
||
# HTTP helpers (requests 优先, 无则 urllib)
|
||
# -------------------------
|
||
def _http_post_json(self, url: str, payload: Dict[str, Any]) -> Dict[str, Any]:
|
||
try:
|
||
import requests # type: ignore
|
||
except Exception:
|
||
requests = None # type: ignore
|
||
|
||
data = json.dumps(payload).encode("utf-8")
|
||
if requests is not None:
|
||
r = requests.post(url, json=payload, timeout=self.config.timeout_s)
|
||
r.raise_for_status()
|
||
return r.json()
|
||
|
||
import urllib.request
|
||
|
||
req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
|
||
with urllib.request.urlopen(req, timeout=self.config.timeout_s) as resp:
|
||
raw = resp.read().decode("utf-8", errors="replace")
|
||
return json.loads(raw)
|
||
|
||
def _http_get_json(self, url: str) -> Any:
|
||
try:
|
||
import requests # type: ignore
|
||
except Exception:
|
||
requests = None # type: ignore
|
||
|
||
if requests is not None:
|
||
r = requests.get(url, timeout=self.config.timeout_s)
|
||
r.raise_for_status()
|
||
return r.json()
|
||
|
||
import urllib.request
|
||
|
||
with urllib.request.urlopen(url, timeout=self.config.timeout_s) as resp:
|
||
raw = resp.read().decode("utf-8", errors="replace")
|
||
return json.loads(raw)
|
||
|
||
def _http_get_bytes(self, url: str) -> bytes:
|
||
try:
|
||
import requests # type: ignore
|
||
except Exception:
|
||
requests = None # type: ignore
|
||
|
||
if requests is not None:
|
||
r = requests.get(url, timeout=self.config.timeout_s)
|
||
r.raise_for_status()
|
||
return r.content
|
||
|
||
import urllib.request
|
||
|
||
with urllib.request.urlopen(url, timeout=self.config.timeout_s) as resp:
|
||
return resp.read()
|
||
|
||
# -------------------------
|
||
# Core workflow injection
|
||
# -------------------------
|
||
def _inject(self, prompt_json: Dict[str, Any]) -> Dict[str, Any]:
|
||
wf = deepcopy(self._workflow_template)
|
||
|
||
# 模型文件：优先使用参数指定，其次使用配置，最后使用模板默认值
|
||
clip_name = prompt_json.get("clip_name") or self.config.clip_name
|
||
unet_name = prompt_json.get("unet_name") or self.config.unet_name
|
||
vae_name = prompt_json.get("vae_name") or self.config.vae_name
|
||
|
||
wf["45"]["inputs"]["clip_name"] = str(clip_name)
|
||
wf["44"]["inputs"]["unet_name"] = str(unet_name)
|
||
wf["15"]["inputs"]["vae_name"] = str(vae_name)
|
||
|
||
# 可选：LoRA 注入（仅 UNET）
|
||
self._inject_loras(wf, prompt_json.get("loras"))
|
||
|
||
# 文本
|
||
positive = (prompt_json.get("positive") or "").strip()
|
||
if not positive:
|
||
positive = build_anima_positive_text(prompt_json)
|
||
negative = (prompt_json.get("neg") or prompt_json.get("negative") or "").strip()
|
||
|
||
wf["11"]["inputs"]["text"] = positive
|
||
wf["12"]["inputs"]["text"] = negative
|
||
|
||
# 分辨率
|
||
width = prompt_json.get("width")
|
||
height = prompt_json.get("height")
|
||
aspect_ratio = (prompt_json.get("aspect_ratio") or "").strip()
|
||
round_to = int(prompt_json.get("round_to") or self.config.round_to)
|
||
|
||
if (width is None or height is None) and aspect_ratio:
|
||
# 仅提供 aspect_ratio 时自动计算
|
||
w, h = estimate_size_from_ratio(
|
||
aspect_ratio=aspect_ratio,
|
||
target_megapixels=float(prompt_json.get("target_megapixels") or self.config.target_megapixels),
|
||
round_to=round_to,
|
||
)
|
||
width, height = w, h
|
||
elif width is not None and height is not None:
|
||
# 用户直接指定了 width/height，也需要对齐到 round_to 的倍数
|
||
# 避免 "should be divisible by spatial_patch_size" 错误
|
||
width = align_dimension(width, round_to)
|
||
height = align_dimension(height, round_to)
|
||
|
||
if width is None or height is None:
|
||
# 默认方形 1MP（1024 是 16 的倍数）
|
||
width, height = 1024, 1024
|
||
|
||
wf["28"]["inputs"]["width"] = int(width)
|
||
wf["28"]["inputs"]["height"] = int(height)
|
||
wf["28"]["inputs"]["batch_size"] = int(prompt_json.get("batch_size") or 1)
|
||
|
||
# 采样参数
|
||
seed = prompt_json.get("seed")
|
||
if seed is None:
|
||
seed = int.from_bytes(uuid.uuid4().bytes[:4], "big", signed=False)
|
||
wf["19"]["inputs"]["seed"] = int(seed)
|
||
|
||
wf["19"]["inputs"]["steps"] = int(prompt_json.get("steps") or wf["19"]["inputs"]["steps"])
|
||
wf["19"]["inputs"]["cfg"] = float(prompt_json.get("cfg") or wf["19"]["inputs"]["cfg"])
|
||
wf["19"]["inputs"]["sampler_name"] = str(prompt_json.get("sampler_name") or wf["19"]["inputs"]["sampler_name"])
|
||
wf["19"]["inputs"]["scheduler"] = str(prompt_json.get("scheduler") or wf["19"]["inputs"]["scheduler"])
|
||
wf["19"]["inputs"]["denoise"] = float(prompt_json.get("denoise") or wf["19"]["inputs"]["denoise"])
|
||
|
||
# 文件名前缀
|
||
wf["52"]["inputs"]["filename_prefix"] = str(prompt_json.get("filename_prefix") or wf["52"]["inputs"]["filename_prefix"])
|
||
|
||
return wf
|
||
|
||
# -------------------------
|
||
# Health check
|
||
# -------------------------
|
||
def check_comfyui_health(self) -> Tuple[bool, str]:
|
||
"""
|
||
检查 ComfyUI 是否可访问。
|
||
返回 (is_healthy, message)
|
||
"""
|
||
try:
|
||
url = urljoin(self.config.comfyui_url.rstrip("/") + "/", "system_stats")
|
||
self._http_get_json(url)
|
||
return True, f"ComfyUI 运行正常 ({self.config.comfyui_url})"
|
||
except Exception as e:
|
||
error_msg = str(e)
|
||
# 提供友好的错误提示
|
||
if "Connection refused" in error_msg or "连接" in error_msg:
|
||
return False, (
|
||
f"无法连接到 ComfyUI ({self.config.comfyui_url})\n"
|
||
f"请确认：\n"
|
||
f" 1. ComfyUI 已启动\n"
|
||
f" 2. 地址和端口正确（可通过 COMFYUI_URL 环境变量修改）\n"
|
||
f" 3. 防火墙未阻止连接"
|
||
)
|
||
elif "timeout" in error_msg.lower() or "超时" in error_msg:
|
||
return False, (
|
||
f"连接 ComfyUI 超时 ({self.config.comfyui_url})\n"
|
||
f"可能原因：网络延迟、ComfyUI 负载过高"
|
||
)
|
||
else:
|
||
return False, f"ComfyUI 连接错误: {error_msg}"
|
||
|
||
# -------------------------
|
||
# ComfyUI execution
|
||
# -------------------------
|
||
def queue_prompt(self, prompt: Dict[str, Any]) -> str:
|
||
url = urljoin(self.config.comfyui_url.rstrip("/") + "/", "prompt")
|
||
payload = {"prompt": prompt, "client_id": self._client_id}
|
||
|
||
try:
|
||
resp = self._http_post_json(url, payload)
|
||
except Exception as e:
|
||
# 先检查 ComfyUI 是否可访问
|
||
is_healthy, health_msg = self.check_comfyui_health()
|
||
if not is_healthy:
|
||
raise RuntimeError(health_msg) from e
|
||
raise
|
||
|
||
prompt_id = str(resp.get("prompt_id") or "")
|
||
if not prompt_id:
|
||
# 检查是否有错误信息
|
||
error = resp.get("error") or resp.get("node_errors")
|
||
if error:
|
||
raise RuntimeError(f"ComfyUI 执行错误：{error}")
|
||
raise RuntimeError(f"ComfyUI /prompt 返回异常：{resp}")
|
||
return prompt_id
|
||
|
||
def wait_history(self, prompt_id: str) -> Dict[str, Any]:
|
||
url = urljoin(self.config.comfyui_url.rstrip("/") + "/", f"history/{prompt_id}")
|
||
deadline = time.time() + float(self.config.timeout_s)
|
||
last = None
|
||
while time.time() < deadline:
|
||
data = self._http_get_json(url)
|
||
last = data
|
||
if isinstance(data, dict) and prompt_id in data:
|
||
return data[prompt_id]
|
||
time.sleep(float(self.config.poll_interval_s))
|
||
raise TimeoutError(f"等待 ComfyUI 生成超时：prompt_id={prompt_id}, last={last}")
|
||
|
||
def _extract_images(self, prompt_id: str, history_item: Dict[str, Any]) -> List[GeneratedImage]:
|
||
outputs = history_item.get("outputs") or {}
|
||
images: List[GeneratedImage] = []
|
||
for node_id, node_out in outputs.items():
|
||
if not isinstance(node_out, dict):
|
||
continue
|
||
for im in node_out.get("images") or []:
|
||
filename = str(im.get("filename") or "")
|
||
subfolder = str(im.get("subfolder") or "")
|
||
folder_type = str(im.get("type") or "output")
|
||
if not filename:
|
||
continue
|
||
qs = urlencode({"filename": filename, "subfolder": subfolder, "type": folder_type})
|
||
view_url = urljoin(self.config.comfyui_url.rstrip("/") + "/", f"view?{qs}")
|
||
images.append(
|
||
GeneratedImage(
|
||
filename=filename,
|
||
subfolder=subfolder,
|
||
folder_type=folder_type,
|
||
view_url=view_url,
|
||
saved_path=None,
|
||
)
|
||
)
|
||
return images
|
||
|
||
def _download_images(self, images: List[GeneratedImage]) -> List[GeneratedImage]:
|
||
"""下载图片并保存到本地，同时保留原始 bytes 用于 base64 编码"""
|
||
out_dir = Path(self.config.output_dir)
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
downloaded: List[GeneratedImage] = []
|
||
for im in images:
|
||
# 总是下载内容（用于 base64）
|
||
content = self._http_get_bytes(im.view_url)
|
||
|
||
saved_path = None
|
||
if self.config.download_images:
|
||
# 复刻 ComfyUI 的 subfolder 结构（可选）
|
||
sub_dir = out_dir / (im.subfolder or "")
|
||
sub_dir.mkdir(parents=True, exist_ok=True)
|
||
dst = sub_dir / im.filename
|
||
dst.write_bytes(content)
|
||
saved_path = str(dst)
|
||
|
||
downloaded.append(
|
||
GeneratedImage(
|
||
filename=im.filename,
|
||
subfolder=im.subfolder,
|
||
folder_type=im.folder_type,
|
||
view_url=im.view_url,
|
||
saved_path=saved_path,
|
||
content=content,
|
||
)
|
||
)
|
||
return downloaded
|
||
|
||
def _get_mime_type(self, filename: str) -> str:
|
||
"""根据文件名推断 MIME 类型"""
|
||
ext = Path(filename).suffix.lower()
|
||
return {
|
||
".png": "image/png",
|
||
".jpg": "image/jpeg",
|
||
".jpeg": "image/jpeg",
|
||
".webp": "image/webp",
|
||
".gif": "image/gif",
|
||
}.get(ext, "image/png")
|
||
|
||
def check_models(self) -> Tuple[bool, str]:
|
||
"""
|
||
检查模型文件是否存在（如果配置了 COMFYUI_MODELS_DIR）。
|
||
返回 (is_ok, message)
|
||
"""
|
||
if not self.config.check_models or not self.config.comfyui_models_dir:
|
||
return True, "模型检查已跳过（未配置 COMFYUI_MODELS_DIR）"
|
||
|
||
all_exist, missing = self.config.check_models_exist()
|
||
if all_exist:
|
||
return True, "所有模型文件已就绪"
|
||
|
||
missing_str = "\n".join(f" - {m}" for m in missing)
|
||
return False, (
|
||
f"缺少以下模型文件：\n{missing_str}\n\n"
|
||
f"请从 HuggingFace 下载：https://huggingface.co/circlestone-labs/Anima\n"
|
||
f"并放置到 ComfyUI/models/ 对应子目录"
|
||
)
|
||
|
||
def generate(self, prompt_json: Dict[str, Any]) -> Dict[str, Any]:
|
||
"""
|
||
输入结构化 JSON，执行生成。
|
||
|
||
返回：
|
||
- prompt_id
|
||
- positive / negative（最终发送给 ComfyUI 的文本）
|
||
- width / height
|
||
- images: [{filename, url, file_path, base64, mime_type, markdown}]
|
||
"""
|
||
# 预检查：模型文件
|
||
models_ok, models_msg = self.check_models()
|
||
if not models_ok:
|
||
raise RuntimeError(models_msg)
|
||
|
||
prompt = self._inject(prompt_json)
|
||
prompt_id = self.queue_prompt(prompt)
|
||
history_item = self.wait_history(prompt_id)
|
||
images = self._extract_images(prompt_id, history_item)
|
||
images = self._download_images(images)
|
||
|
||
# 构建图片信息（包含多种格式）
|
||
images_data = []
|
||
for im in images:
|
||
mime_type = self._get_mime_type(im.filename)
|
||
b64 = base64.b64encode(im.content).decode("ascii") if im.content else None
|
||
|
||
img_info = {
|
||
"filename": im.filename,
|
||
"subfolder": im.subfolder,
|
||
"type": im.folder_type,
|
||
# URL 格式
|
||
"url": im.view_url,
|
||
"view_url": im.view_url, # 兼容旧字段
|
||
# 本地路径
|
||
"file_path": im.saved_path,
|
||
"saved_path": im.saved_path, # 兼容旧字段
|
||
# Base64 格式（用于 MCP / Gemini / 嵌入）
|
||
"base64": b64,
|
||
"mime_type": mime_type,
|
||
# Data URL（可直接用于 <img src> 或 markdown）
|
||
"data_url": f"data:{mime_type};base64,{b64}" if b64 else None,
|
||
# Markdown 格式（AI 可直接输出）
|
||
"markdown": f"",
|
||
}
|
||
images_data.append(img_info)
|
||
|
||
# 回显最终参数（便于调试）
|
||
actual_seed = int(prompt["19"]["inputs"]["seed"])
|
||
actual_width = int(prompt["28"]["inputs"]["width"])
|
||
actual_height = int(prompt["28"]["inputs"]["height"])
|
||
|
||
result = {
|
||
"success": True,
|
||
"prompt_id": prompt_id,
|
||
"positive": prompt["11"]["inputs"]["text"],
|
||
"negative": prompt["12"]["inputs"]["text"],
|
||
"seed": actual_seed,
|
||
"width": actual_width,
|
||
"height": actual_height,
|
||
"images": images_data,
|
||
}
|
||
|
||
# 记录到历史
|
||
record = self.history.add(
|
||
params=prompt_json,
|
||
positive_text=result["positive"],
|
||
negative_text=result["negative"],
|
||
prompt_id=prompt_id,
|
||
seed=actual_seed,
|
||
width=actual_width,
|
||
height=actual_height,
|
||
)
|
||
result["history_id"] = record.id
|
||
|
||
return result
|