wehub-resource-sync
154 lines
4.5 KiB
Python
154 lines
4.5 KiB
Python
"""MetaStyle T8 节点核心逻辑
|
||
|
||
输入:
|
||
- style_key (string): 选中的风格 key，前端 widget 写入
|
||
- output_size (combo): 输出尺寸 original / 512 / 1024
|
||
- seed (int): 预留
|
||
输出:
|
||
- IMAGE 风格代表图 (BHWC float32 0~1)
|
||
- style_name 风格全名
|
||
- category 大类（中文）
|
||
- sample_prompt 示例提示词（可直接接 CLIPTextEncode）
|
||
- sample_id 数据集原始 id
|
||
- meta_json 全量元数据 JSON 字符串
|
||
"""
|
||
import json
|
||
import os
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import torch
|
||
from PIL import Image
|
||
|
||
ROOT = Path(__file__).resolve().parent
|
||
DATA = ROOT / "data"
|
||
|
||
_meta_cache = None
|
||
_catalog_cache = None
|
||
_lite_cache = None
|
||
|
||
|
||
def load_meta() -> dict:
|
||
global _meta_cache
|
||
if _meta_cache is None:
|
||
p = DATA / "styles_meta.json"
|
||
_meta_cache = json.loads(p.read_text(encoding="utf-8")) if p.exists() else {}
|
||
return _meta_cache
|
||
|
||
|
||
def load_catalog() -> dict:
|
||
global _catalog_cache
|
||
if _catalog_cache is None:
|
||
p = DATA / "catalog.json"
|
||
_catalog_cache = json.loads(p.read_text(encoding="utf-8")) if p.exists() else {}
|
||
return _catalog_cache
|
||
|
||
|
||
def load_lite() -> list:
|
||
global _lite_cache
|
||
if _lite_cache is None:
|
||
p = DATA / "styles_lite.json"
|
||
_lite_cache = json.loads(p.read_text(encoding="utf-8")) if p.exists() else []
|
||
return _lite_cache
|
||
|
||
|
||
def pil_to_tensor(im: Image.Image) -> torch.Tensor:
|
||
"""PIL -> ComfyUI 标准 IMAGE tensor: shape (1, H, W, 3), float32 0~1"""
|
||
if im.mode != "RGB":
|
||
im = im.convert("RGB")
|
||
arr = np.asarray(im).astype(np.float32) / 255.0
|
||
return torch.from_numpy(arr)[None, ...]
|
||
|
||
|
||
def empty_image(w: int = 512, h: int = 512) -> torch.Tensor:
|
||
return torch.zeros((1, h, w, 3), dtype=torch.float32)
|
||
|
||
|
||
class MetaStyleT8Picker:
|
||
"""T8 风格选择器：按风格挑选代表图与示例提示词。"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
# 由前端 widget 维护的字符串：选中的 style key
|
||
"style_key": ("STRING", {
|
||
"default": "",
|
||
"multiline": False,
|
||
"metastyle_picker": True, # 给前端 JS 识别用
|
||
}),
|
||
"output_size": (["original", "512", "768", "1024"],
|
||
{"default": "original"}),
|
||
},
|
||
"optional": {
|
||
"seed": ("INT", {
|
||
"default": 0, "min": 0, "max": 0xFFFFFFFF
|
||
}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = (
|
||
"IMAGE", "STRING", "STRING", "STRING", "STRING", "STRING",
|
||
)
|
||
RETURN_NAMES = (
|
||
"image", "style_name", "category",
|
||
"sample_prompt", "sample_id", "meta_json",
|
||
)
|
||
FUNCTION = "pick"
|
||
CATEGORY = "T8/Style"
|
||
|
||
def pick(self, style_key: str, output_size: str, seed: int = 0):
|
||
meta = load_meta()
|
||
style_key = (style_key or "").strip()
|
||
info = meta.get(style_key)
|
||
|
||
if not info:
|
||
# 未选中或 key 无效 → 返回空图，提示文字
|
||
return (
|
||
empty_image(),
|
||
style_key,
|
||
"",
|
||
"",
|
||
"",
|
||
json.dumps({"error": "style_key not found", "key": style_key},
|
||
ensure_ascii=False),
|
||
)
|
||
|
||
img_path = ROOT / "data" / info["image"]
|
||
try:
|
||
with Image.open(img_path) as im:
|
||
im.load()
|
||
if output_size != "original":
|
||
target = int(output_size)
|
||
im = im.copy()
|
||
im.thumbnail((target, target), Image.LANCZOS)
|
||
tensor = pil_to_tensor(im)
|
||
except Exception as e:
|
||
return (
|
||
empty_image(),
|
||
style_key,
|
||
info.get("category", ""),
|
||
info.get("sample_content", ""),
|
||
info.get("sample_id", ""),
|
||
json.dumps({"error": str(e), "key": style_key},
|
||
ensure_ascii=False),
|
||
)
|
||
|
||
return (
|
||
tensor,
|
||
info.get("style_full", style_key),
|
||
info.get("category", ""),
|
||
info.get("sample_content", ""),
|
||
info.get("sample_id", ""),
|
||
json.dumps(info, ensure_ascii=False),
|
||
)
|
||
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"MetaStyleT8Picker": MetaStyleT8Picker,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"MetaStyleT8Picker": "MetaStyle T8 风格选择器",
|
||
}
|