249 lines
9.3 KiB
Python
249 lines
9.3 KiB
Python
"""权重分配、提示词构建/解析与去重键。
|
|
|
|
纯逻辑模块，不涉及文件 IO，便于单独测试。
|
|
"""
|
|
|
|
import random
|
|
import re
|
|
from typing import Iterable, List, Optional, Sequence, Tuple
|
|
|
|
from .constants import WEIGHT_SCALE
|
|
|
|
__all__ = [
|
|
"DEDUP_MODES",
|
|
"WEIGHT_CURVES",
|
|
"generate_weights",
|
|
"build_prompt",
|
|
"parse_prompt",
|
|
"parse_weight_list",
|
|
"normalize_dedup_mode",
|
|
"normalize_weight_curve",
|
|
"escape_key_part",
|
|
"dedup_key",
|
|
]
|
|
|
|
# 去重模式：只有 none 之外的模式才需要读取已测记录
|
|
DEDUP_MODES = ("none", "full_prompt", "artist_set", "artist_list")
|
|
|
|
# 权重分配曲线（决定各艺术家权重是平均分配还是主次分明）：
|
|
# - flat 现状：每位艺术家的随机倾向接近（0.5~1.5），风格平均混合
|
|
# - dominant 主次分明：随机挑一位当主风格，权重明显高于其余
|
|
# - ramp 阶梯：按随机名次递减（第一个最多，依次递减）
|
|
WEIGHT_CURVES = ("flat", "dominant", "ramp")
|
|
|
|
# 解析 "artist:name:0.8" / "(artist:name:0.8)" / "name" 这类片段
|
|
_PROMPT_PART_RE = re.compile(r"^\(*\s*(?:artist:)?\s*(.*?)\s*(?::\s*([0-9]*\.?[0-9]+))?\s*\)*$")
|
|
_KEY_SEP = "|"
|
|
|
|
|
|
def _to_units(value: float) -> int:
|
|
"""把权重值换算成步长为 WEIGHT_STEP 的整数单位（四舍五入到 0.1）。"""
|
|
return int(round(float(value) * WEIGHT_SCALE))
|
|
|
|
|
|
def _from_units(units: int) -> float:
|
|
"""把整数单位还原为保留一位小数的权重值。"""
|
|
return round(units / WEIGHT_SCALE, 1)
|
|
|
|
|
|
def _curve_desires(count: int, curve: str, rng: random.Random) -> List[float]:
|
|
"""按分配曲线生成每个槽位的“倾向值”（相对份额，只有比例有意义）。
|
|
|
|
- ``flat``：全部在 0.5~1.5 之间随机，分配结果接近平均
|
|
- ``dominant``：随机挑一位给高倾向，其余给低倾向，形成明显主次
|
|
- ``ramp``：把槽位随机打乱后按名次递减（1, 1/2^0.7, 1/3^0.7 …），形成阶梯
|
|
"""
|
|
if count <= 0:
|
|
return []
|
|
if curve == "dominant":
|
|
main = rng.randrange(count)
|
|
return [rng.uniform(2.2, 3.2) if index == main else rng.uniform(0.15, 0.6) for index in range(count)]
|
|
if curve == "ramp":
|
|
order = list(range(count))
|
|
rng.shuffle(order)
|
|
desires = [0.0] * count
|
|
for rank, index in enumerate(order):
|
|
desires[index] = 1.0 / (rank + 1) ** 0.7
|
|
return desires
|
|
return [rng.uniform(0.5, 1.5) for _ in range(count)]
|
|
|
|
|
|
def generate_weights(
|
|
count: int,
|
|
weight_min: float,
|
|
weight_max: float,
|
|
weight_total: float,
|
|
rng: random.Random,
|
|
curve: Optional[str] = None,
|
|
) -> List[float]:
|
|
"""生成满足约束的随机权重列表。
|
|
|
|
约束（与节点界面描述一致）：
|
|
1. 每个权重 ∈ [weight_min, weight_max]（不会越界；weight_max < weight_min 时自动交换）
|
|
2. 权重总和精确等于 weight_total（超出可行区间时夹到边界）
|
|
3. 每个权重保留一位小数，且分配结果带随机性
|
|
4. 权重恒为非负：传入负数（API 直传可能出现）会被夹到 0
|
|
|
|
参数:
|
|
curve: 分配曲线，见 :data:`WEIGHT_CURVES`（非法值回退为 flat）。
|
|
它只影响“谁分得多”，不会破坏上面 4 条约束。
|
|
|
|
实现说明：所有运算都在“0.1 的整数倍”单位上进行，避免浮点累计误差；
|
|
先按随机倾向值做近似平均的分配，再用最大余数法与容量回填保证总和精确，
|
|
因此不会再出现旧实现中“为凑总和而越过 weight_max”的问题。
|
|
"""
|
|
n = int(count)
|
|
if n <= 0:
|
|
return []
|
|
|
|
low_units = _to_units(weight_min)
|
|
high_units = _to_units(weight_max)
|
|
if high_units < low_units: # 容忍参数写反的情况
|
|
low_units, high_units = high_units, low_units
|
|
# 边界保护：负数权重在提示词里无意义，会污染总和与条件编码
|
|
low_units = max(0, low_units)
|
|
high_units = max(0, high_units)
|
|
|
|
total_units = _to_units(weight_total)
|
|
# 夹到可行区间：n * min <= total <= n * max
|
|
total_units = max(n * low_units, min(total_units, n * high_units))
|
|
extra = total_units - n * low_units
|
|
if extra <= 0:
|
|
return [_from_units(low_units)] * n
|
|
|
|
caps = [high_units - low_units] * n # 每个槽位在最小值之上还能分配多少
|
|
if sum(caps) <= 0: # min == max，所有槽位只能取同一个值
|
|
return [_from_units(low_units)] * n
|
|
|
|
# 1) 倾向值决定各槽位的份额；曲线决定分布是平均还是主次分明
|
|
desire = _curve_desires(n, normalize_weight_curve(curve), rng)
|
|
desire_sum = sum(desire)
|
|
shares = [extra * value / desire_sum for value in desire]
|
|
|
|
allocation = [min(int(share), cap) for share, cap in zip(shares, caps)]
|
|
remain = extra - sum(allocation)
|
|
|
|
# 2) 余数按小数部分从大到小补齐（最大余数法），保证总和精确
|
|
order = sorted(
|
|
range(n),
|
|
key=lambda index: (shares[index] - int(shares[index]), rng.random()),
|
|
reverse=True,
|
|
)
|
|
for index in order:
|
|
if remain <= 0:
|
|
break
|
|
if allocation[index] < caps[index]:
|
|
allocation[index] += 1
|
|
remain -= 1
|
|
|
|
# 3) 仍有余量时逐个回填（sum(caps) >= extra 恒成立，一定能填满）
|
|
while remain > 0:
|
|
progressed = False
|
|
for index in range(n):
|
|
if remain <= 0:
|
|
break
|
|
if allocation[index] < caps[index]:
|
|
allocation[index] += 1
|
|
remain -= 1
|
|
progressed = True
|
|
if not progressed:
|
|
break
|
|
|
|
return [_from_units(low_units + value) for value in allocation]
|
|
|
|
|
|
def build_prompt(artists: Iterable[str], weights: Iterable[float], prefix: str = "") -> str:
|
|
"""把艺术家名与权重构建为提示词字符串。
|
|
|
|
权重为 1.0 时直接输出名称（不含括号与权重），例如 ``artist:a,(artist:b:0.7),``。
|
|
|
|
末尾**总是补一个逗号**（即使只抽到 1 个艺术家）：这样把画师提示词与角色节点的
|
|
多行输出直接拼接时不会粘在一起。:func:`parse_prompt` 会跳过结尾的空段，
|
|
因此「构建 -> 解析」的往返结果不受影响。
|
|
"""
|
|
parts: List[str] = []
|
|
for artist, weight in zip(artists, weights):
|
|
value = round(float(weight), 1)
|
|
name = f"{prefix}{artist}"
|
|
parts.append(name if value == 1.0 else f"({name}:{value})")
|
|
if not parts:
|
|
return ""
|
|
return ",".join(parts) + ","
|
|
|
|
|
|
def parse_prompt(prompt: str) -> Tuple[List[str], List[float]]:
|
|
"""从提示词字符串反向解析出艺术家名与权重列表。
|
|
|
|
支持 ``artist:a,(artist:b:0.7)`` 或 ``a,(b:0.7)`` 两种形式。
|
|
"""
|
|
artists: List[str] = []
|
|
weights: List[float] = []
|
|
if not prompt:
|
|
return artists, weights
|
|
for part in str(prompt).split(","):
|
|
part = part.strip()
|
|
if not part:
|
|
continue
|
|
match = _PROMPT_PART_RE.match(part)
|
|
if match is None:
|
|
name, weight = part.strip("() "), 1.0
|
|
else:
|
|
name = (match.group(1) or "").strip()
|
|
weight = float(match.group(2)) if match.group(2) else 1.0
|
|
if name:
|
|
artists.append(name)
|
|
weights.append(round(weight, 1))
|
|
return artists, weights
|
|
|
|
|
|
def parse_weight_list(text: str) -> List[float]:
|
|
"""解析逗号分隔的权重字符串（已测 CSV 第二列），无法解析的项会被跳过。"""
|
|
weights: List[float] = []
|
|
for item in str(text or "").split(","):
|
|
item = item.strip()
|
|
if not item:
|
|
continue
|
|
try:
|
|
weights.append(round(float(item), 1))
|
|
except ValueError:
|
|
continue
|
|
return weights
|
|
|
|
|
|
def normalize_dedup_mode(mode: Optional[str]) -> str:
|
|
"""规范化去重模式，非法值回退为 none。"""
|
|
value = str(mode or "none").strip().lower()
|
|
return value if value in DEDUP_MODES else "none"
|
|
|
|
|
|
def normalize_weight_curve(curve: Optional[str]) -> str:
|
|
"""规范化权重分配曲线，非法值回退为 flat（保持历史行为）。"""
|
|
value = str(curve or "flat").strip().lower()
|
|
return value if value in WEIGHT_CURVES else "flat"
|
|
|
|
|
|
def escape_key_part(value: str) -> str:
|
|
"""转义键分隔符：把 | 与 : 变成重复字符（可逆的重复表示法）。
|
|
|
|
保持纯字符串变换（不依赖任何外部状态），便于单独测试与复用。
|
|
"""
|
|
return str(value).replace(_KEY_SEP, _KEY_SEP * 2).replace(":", "::")
|
|
|
|
|
|
def dedup_key(mode: str, artists: Sequence[str], weights: Sequence[float]) -> str:
|
|
"""按去重模式计算组合的唯一键；mode 为 none 时返回空字符串。
|
|
|
|
名称中的分隔符会被转义，因此 ``["a:b"]`` 与 ``["a", "b"]`` 不会拼出同一个键。
|
|
旧实现直接拼接，艺术家名含 ``:`` / ``|`` 时会把不同组合判成同一个，
|
|
导致这些组合永远不被输出（去重假阳性）。
|
|
"""
|
|
if mode == "full_prompt":
|
|
return _KEY_SEP.join(
|
|
f"{escape_key_part(artist)}:{round(float(weight), 1)}" for artist, weight in zip(artists, weights)
|
|
)
|
|
if mode == "artist_set":
|
|
return _KEY_SEP.join(sorted(escape_key_part(name) for name in set(artists)))
|
|
if mode == "artist_list":
|
|
return _KEY_SEP.join(sorted(escape_key_part(name) for name in artists))
|
|
return ""
|