Add files via upload
This commit is contained in:
Ruinan
@@ -0,0 +1,292 @@
|
||||
import json
|
||||
import os
|
||||
import hashlib
|
||||
import time
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
"""
|
||||
最简批量帧文本编码（仅正向），增强兼容性：
|
||||
- 解决：不同 prompt 生成的 cond 长度不一致导致 torch.cat 报错 -> 统一 padding
|
||||
- 解决：某些 clip.encode_from_tokens 在 return_pooled=True 时返回 (cond, None) 或根本不支持 pooled 的情况
|
||||
- 若 pooled 全部为 None，则不放 pooled_output 或放占位空字符串
|
||||
"""
|
||||
|
||||
# ComfyUI 节点注册字典（在文件末尾进行 update 之前，必须先定义）
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
def parse_frame_prompts(raw_text: str):
|
||||
"""
|
||||
解析输入为 {frame_index: prompt} 的升序字典。
|
||||
允许缺少最外层 {}；允许末尾多余逗号。
|
||||
"""
|
||||
txt = raw_text.strip()
|
||||
if not txt:
|
||||
return {}
|
||||
if not txt.startswith("{"):
|
||||
txt = "{" + txt
|
||||
if not txt.endswith("}"):
|
||||
txt = txt + "}"
|
||||
txt = txt.replace(",}", "}")
|
||||
data = json.loads(txt)
|
||||
out = {}
|
||||
for k, v in data.items():
|
||||
try:
|
||||
idx = int(k)
|
||||
except:
|
||||
continue
|
||||
# 非数字 key 忽略
|
||||
out[idx] = str(v)
|
||||
return dict(sorted(out.items(), key=lambda x: x[0]))
|
||||
|
||||
|
||||
def encode_prompts_sequential(clip, ordered_items):
|
||||
"""
|
||||
逐条编码；对 cond 做长度统一；pooled 可能为 None。
|
||||
ordered_items: List[(frame_index, prompt)]
|
||||
返回: cond_final, pooled_final(或 None), has_pooled(bool)
|
||||
"""
|
||||
cond_list = []
|
||||
pooled_list = []
|
||||
pooled_available = True # 先假设有 pooled
|
||||
max_tokens = 0
|
||||
|
||||
with torch.no_grad():
|
||||
for frame_idx, prompt in ordered_items:
|
||||
tokens = clip.tokenize(prompt)
|
||||
# 兼容：有的实现如果 return_pooled=True 仍返回单值
|
||||
enc = clip.encode_from_tokens(tokens, return_pooled=True)
|
||||
if isinstance(enc, (list, tuple)) and len(enc) == 2:
|
||||
cond, pooled = enc
|
||||
else:
|
||||
cond, pooled = enc, None
|
||||
|
||||
if cond.dim() != 3:
|
||||
# 期望 shape: [B(=1), T, C]; 若是 [T, C] 补一个 batch 维
|
||||
if cond.dim() == 2:
|
||||
cond = cond.unsqueeze(0)
|
||||
else:
|
||||
raise ValueError(f"Unexpected cond tensor shape: {cond.shape}")
|
||||
|
||||
if cond.shape[1] > max_tokens:
|
||||
max_tokens = cond.shape[1]
|
||||
|
||||
cond_list.append(cond)
|
||||
|
||||
if pooled is None:
|
||||
pooled_available = False
|
||||
pooled_list.append(None)
|
||||
else:
|
||||
# 期望 [B(=1), D] 或 [D]; 若 [D] -> [1,D]
|
||||
if pooled.dim() == 1:
|
||||
pooled = pooled.unsqueeze(0)
|
||||
pooled_list.append(pooled)
|
||||
|
||||
if not cond_list:
|
||||
dummy = torch.zeros(1, 1)
|
||||
return dummy, None, False
|
||||
|
||||
# 对 cond 做 padding
|
||||
padded = []
|
||||
for c in cond_list:
|
||||
if c.shape[1] < max_tokens:
|
||||
pad_len = max_tokens - c.shape[1]
|
||||
# shape [1, T, C]，pad 第二维 (tokens)：(last_dim_left, last_dim_right, second_dim_left, second_dim_right)
|
||||
c = F.pad(c, (0, 0, 0, pad_len))
|
||||
padded.append(c)
|
||||
cond_final = torch.cat(padded, dim=0) # [N, max_tokens, C]
|
||||
|
||||
pooled_final = None
|
||||
if pooled_available:
|
||||
# 校验 pooled 维度并 cat
|
||||
processed = []
|
||||
for p in pooled_list:
|
||||
if p is None:
|
||||
pooled_available = False
|
||||
break
|
||||
if p.dim() == 1:
|
||||
p = p.unsqueeze(0)
|
||||
processed.append(p)
|
||||
if pooled_available and processed:
|
||||
try:
|
||||
pooled_final = torch.cat(processed, dim=0) # [N, D]
|
||||
except Exception:
|
||||
# 如果维度不一致，放弃 pooled
|
||||
pooled_final = None
|
||||
pooled_available = False
|
||||
|
||||
return cond_final, pooled_final, pooled_available
|
||||
|
||||
|
||||
## 删除了 BatchPrompt 节点，保留 PromptQueue
|
||||
|
||||
|
||||
def _read_lines_from_source(multiline_text: str, file_path: str = "", use_file: bool = False):
|
||||
"""
|
||||
从多行文本或文件读取，每行一个提示词；自动去除空行和首尾空白。
|
||||
"""
|
||||
lines = []
|
||||
if use_file and file_path:
|
||||
try:
|
||||
with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
|
||||
lines = f.readlines()
|
||||
except Exception:
|
||||
# 回退到多行文本
|
||||
lines = multiline_text.splitlines()
|
||||
else:
|
||||
lines = multiline_text.splitlines()
|
||||
|
||||
cleaned = []
|
||||
for raw in lines:
|
||||
s = str(raw).strip()
|
||||
if not s:
|
||||
continue
|
||||
cleaned.append(s)
|
||||
return cleaned
|
||||
|
||||
|
||||
def _apply_template_to_lines(lines, template: str):
|
||||
"""
|
||||
将每行替换进模板。若模板为空或不含 {p}，则直接返回行本身。
|
||||
"""
|
||||
if not template or "{p}" not in template:
|
||||
return lines
|
||||
return [template.replace("{p}", line) for line in lines]
|
||||
|
||||
|
||||
class PromptQueue:
|
||||
"""
|
||||
将多行提示词视为队列，或从 txt 文件读取（逐行）。
|
||||
输出按批次堆叠的 CONDITIONING，便于一次性批量生成。
|
||||
可选模板，将每行填充到模板中的 {prompt} 占位符。
|
||||
"""
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"clip": ("CLIP",),
|
||||
"multiline_text": ("STRING", {"multiline": True, "default": "A cat\nA dog\nA bird"}),
|
||||
"use_file": ("BOOLEAN", {"default": False}),
|
||||
"file_path": ("STRING", {"multiline": False, "default": ""}),
|
||||
"template": ("STRING", {"multiline": False, "default": "{p}"}),
|
||||
"mode": (["all", "incremental"],),
|
||||
"label": ("STRING", {"multiline": False, "default": "Prompt Queue 001"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING",)
|
||||
RETURN_NAMES = ("COND",)
|
||||
FUNCTION = "run"
|
||||
CATEGORY = "SimpleBatch"
|
||||
|
||||
def run(self, clip, multiline_text, use_file, file_path, template, mode="all", label="Prompt Queue 001"):
|
||||
lines = _read_lines_from_source(multiline_text, file_path, use_file)
|
||||
lines = _apply_template_to_lines(lines, template)
|
||||
|
||||
if not lines:
|
||||
dummy = torch.zeros(1, 1)
|
||||
return ([[dummy, {"pooled_output": ""}]],)
|
||||
|
||||
# Mode: incremental consumes one line per call using a persistent counter per label
|
||||
if mode == "incremental":
|
||||
idx = _pq_get_next_index(label, _pq_total=len(lines), source_descriptor=_pq_source_descriptor(use_file, file_path, multiline_text))
|
||||
line = lines[idx]
|
||||
ordered = [(0, line)]
|
||||
cond, pooled, has_pooled = encode_prompts_sequential(clip, ordered)
|
||||
meta = {"pooled_output": pooled if has_pooled and pooled is not None else ""}
|
||||
return ([[cond, meta]],)
|
||||
|
||||
# Default: all lines batched together
|
||||
ordered = list(enumerate(lines))
|
||||
cond, pooled, has_pooled = encode_prompts_sequential(clip, ordered)
|
||||
meta = {"pooled_output": pooled if has_pooled and pooled is not None else ""}
|
||||
return ([[cond, meta]],)
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, clip, multiline_text, use_file, file_path, template, mode="all", label="Prompt Queue 001"):
|
||||
# Ensure incremental mode re-executes each time to advance the queue
|
||||
if mode == "incremental":
|
||||
return str(time.time())
|
||||
# For deterministic modes, hash inputs so caching works as expected
|
||||
try:
|
||||
desc = _pq_source_descriptor(bool(use_file), str(file_path), str(multiline_text))
|
||||
key = f"{desc}|{template}"
|
||||
return hashlib.sha1(key.encode("utf-8", errors="ignore")).hexdigest()
|
||||
except Exception:
|
||||
return ""
|
||||
|
||||
|
||||
# --- Persistent queue state (per label) ---
|
||||
class _JsonKVStore:
|
||||
def __init__(self, file_path: str):
|
||||
self.file_path = file_path
|
||||
self.data = {}
|
||||
self._load()
|
||||
|
||||
def _load(self):
|
||||
try:
|
||||
if os.path.exists(self.file_path):
|
||||
with open(self.file_path, "r", encoding="utf-8") as f:
|
||||
self.data = json.load(f)
|
||||
else:
|
||||
self.data = {}
|
||||
except Exception:
|
||||
self.data = {}
|
||||
|
||||
def _save(self):
|
||||
try:
|
||||
os.makedirs(os.path.dirname(self.file_path), exist_ok=True)
|
||||
with open(self.file_path, "w", encoding="utf-8") as f:
|
||||
json.dump(self.data, f, indent=2)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def get(self, category: str, key: str, default=None):
|
||||
return self.data.get(category, {}).get(key, default)
|
||||
|
||||
def put(self, category: str, key: str, value) -> None:
|
||||
if category not in self.data:
|
||||
self.data[category] = {}
|
||||
self.data[category][key] = value
|
||||
self._save()
|
||||
|
||||
|
||||
_pq_state_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "_prompt_queue_state")
|
||||
_pq_store = _JsonKVStore(os.path.join(_pq_state_dir, "state.json"))
|
||||
|
||||
|
||||
def _pq_source_descriptor(use_file: bool, file_path: str, multiline_text: str) -> str:
|
||||
if use_file and file_path:
|
||||
try:
|
||||
return "FILE:" + os.path.abspath(file_path)
|
||||
except Exception:
|
||||
return "FILE:" + file_path
|
||||
# Text hash as descriptor; content changes won't reset index to allow continuing when appending
|
||||
h = hashlib.sha1(multiline_text.encode("utf-8", errors="ignore")).hexdigest()
|
||||
return "TEXT:" + h
|
||||
|
||||
|
||||
def _pq_get_next_index(label: str, _pq_total: int, source_descriptor: str) -> int:
|
||||
if _pq_total <= 0:
|
||||
return 0
|
||||
# Reset only when source type switches between FILE and TEXT, or file path changes
|
||||
stored_src = _pq_store.get("PromptQueue Sources", label, None)
|
||||
if stored_src != source_descriptor:
|
||||
_pq_store.put("PromptQueue Sources", label, source_descriptor)
|
||||
_pq_store.put("PromptQueue Counters", label, 0)
|
||||
|
||||
current = _pq_store.get("PromptQueue Counters", label, 0) or 0
|
||||
next_index = current % _pq_total
|
||||
_pq_store.put("PromptQueue Counters", label, (current + 1) % _pq_total)
|
||||
return next_index
|
||||
|
||||
|
||||
# 注册新节点
|
||||
NODE_CLASS_MAPPINGS.update({
|
||||
"PromptQueue": PromptQueue,
|
||||
})
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update({
|
||||
"PromptQueue": "Prompt Queue",
|
||||
})
|
||||
Reference in New Issue
Block a user