mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-06 14:10:13 -03:00
fix(wildcards): resolve weighted N::value syntax inside wildcard YAML lists (#1039)
This commit is contained in:
@@ -19,7 +19,7 @@ logger = logging.getLogger(__name__)
|
||||
_WILDCARD_PATTERN = re.compile(r"__([\w\s.\-+/*\\]+?)__")
|
||||
_OPTION_PATTERN = re.compile(r"{([^{}]*?)}")
|
||||
_TRIGGER_WORD_PATTERN = re.compile(r"^trigger_words\d+$")
|
||||
_WEIGHTED_OPTION_PATTERN = re.compile(r"^\s*([0-9.]+)::")
|
||||
_WEIGHTED_OPTION_PATTERN = re.compile(r"^\s*-?\d+(\.\d+)?::")
|
||||
_NUMERIC_PATTERN = re.compile(r"^-?\d+(\.\d+)?$")
|
||||
|
||||
|
||||
@@ -390,7 +390,7 @@ class WildcardService:
|
||||
) -> str | None:
|
||||
keyword = _normalize_wildcard_key(raw_key)
|
||||
if keyword in wildcard_dict:
|
||||
return rng.choice(wildcard_dict[keyword])
|
||||
return self._pick_weighted_or_plain(wildcard_dict[keyword], rng)
|
||||
|
||||
if "*" in keyword:
|
||||
regex_pattern = keyword.replace("*", ".*").replace("+", r"\+")
|
||||
@@ -400,7 +400,7 @@ class WildcardService:
|
||||
if compiled.match(key):
|
||||
aggregated.extend(values)
|
||||
if aggregated:
|
||||
return rng.choice(aggregated)
|
||||
return self._pick_weighted_or_plain(aggregated, rng)
|
||||
|
||||
if "/" not in keyword:
|
||||
fallback_keyword = _normalize_wildcard_key(f"*/{keyword}")
|
||||
@@ -409,6 +409,39 @@ class WildcardService:
|
||||
|
||||
return None
|
||||
|
||||
def _pick_weighted_or_plain(
|
||||
self, values: list[str], rng: random.Random
|
||||
) -> str:
|
||||
"""Pick a value from the list, respecting N::weight prefix if present.
|
||||
|
||||
When any value in the list uses the ``N::value`` weighted syntax with a
|
||||
weight different from 1, the pick uses weighted random selection. When
|
||||
no such weighting is present, a plain ``rng.choice`` is used (preserving
|
||||
backward compatibility for unweighted wildcard files).
|
||||
|
||||
In either case the ``N::`` prefix is always stripped from the returned
|
||||
value, matching the behaviour of ``{...}`` option groups.
|
||||
"""
|
||||
# Fast path: skip weighting logic entirely when no :: syntax exists
|
||||
if not any("::" in v for v in values):
|
||||
return rng.choice(values)
|
||||
|
||||
weighted_options: list[tuple[float, str]] = []
|
||||
for value in values:
|
||||
weight = 1.0
|
||||
parts = value.split("::", 1)
|
||||
if len(parts) == 2 and _is_numeric_string(parts[0].strip()):
|
||||
weight = float(parts[0].strip())
|
||||
weighted_options.append((weight, value))
|
||||
|
||||
any_weighted = any(w != 1.0 for w, _ in weighted_options)
|
||||
if any_weighted:
|
||||
picked = self._weighted_choice(weighted_options, rng)
|
||||
else:
|
||||
picked = rng.choice(values)
|
||||
|
||||
return self._strip_weight_prefix(picked)
|
||||
|
||||
|
||||
def is_trigger_words_input(name: str) -> bool:
|
||||
return bool(_TRIGGER_WORD_PATTERN.match(name))
|
||||
|
||||
@@ -139,3 +139,122 @@ def test_contains_dynamic_syntax_detects_wildcards_and_options():
|
||||
assert contains_dynamic_syntax("__flower__") is True
|
||||
assert contains_dynamic_syntax("{red|blue}") is True
|
||||
assert contains_dynamic_syntax("{2$$, $$red|blue|green}") is True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _pick_weighted_or_plain
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_pick_weighted_or_plain_plain_values(monkeypatch, tmp_path):
|
||||
"""Plain values without :: are picked via rng.choice (fast path)."""
|
||||
service, _ = _make_service(monkeypatch, tmp_path)
|
||||
|
||||
import random
|
||||
rng = random.Random(42)
|
||||
|
||||
result = service._pick_weighted_or_plain(["red", "green", "blue"], rng)
|
||||
assert result in {"red", "green", "blue"}
|
||||
assert "::" not in result
|
||||
|
||||
|
||||
def test_pick_weighted_or_plain_deterministic_with_seed(monkeypatch, tmp_path):
|
||||
"""Same seed produces the same result for plain values."""
|
||||
service, _ = _make_service(monkeypatch, tmp_path)
|
||||
|
||||
import random
|
||||
first = service._pick_weighted_or_plain(["a", "b", "c"], random.Random(99))
|
||||
second = service._pick_weighted_or_plain(["a", "b", "c"], random.Random(99))
|
||||
assert first == second
|
||||
|
||||
|
||||
def test_pick_weighted_or_plain_weighted_values(monkeypatch, tmp_path):
|
||||
"""Weighted values use weighted selection and strip the N:: prefix."""
|
||||
service, _ = _make_service(monkeypatch, tmp_path)
|
||||
|
||||
import random
|
||||
values = ["3::apple", "1::banana"]
|
||||
results = {"apple": 0, "banana": 0}
|
||||
for seed in range(4000):
|
||||
result = service._pick_weighted_or_plain(values, random.Random(seed))
|
||||
assert result in results, f"Unexpected result: {result!r}"
|
||||
assert "::" not in result
|
||||
results[result] += 1
|
||||
|
||||
total = results["apple"] + results["banana"]
|
||||
# 3:1 weight → apple ≈ 75%, banana ≈ 25%
|
||||
assert 2700 < results["apple"] < 3300, f"apple count out of range: {results['apple']}"
|
||||
assert 700 < results["banana"] < 1300, f"banana count out of range: {results['banana']}"
|
||||
|
||||
|
||||
def test_pick_weighted_or_plain_weight_one_values(monkeypatch, tmp_path):
|
||||
"""Values with explicit 1:: prefix have prefix stripped but are not weighted."""
|
||||
service, _ = _make_service(monkeypatch, tmp_path)
|
||||
|
||||
import random
|
||||
# All weights are 1.0 → no actual weighting, but :: prefix is stripped
|
||||
values = ["1::foo", "1::bar"]
|
||||
rng = random.Random(42)
|
||||
results = {service._pick_weighted_or_plain(values, rng) for _ in range(200)}
|
||||
assert results == {"foo", "bar"}
|
||||
# Ensure the prefix is always stripped
|
||||
for result in results:
|
||||
assert "::" not in result
|
||||
|
||||
|
||||
def test_pick_weighted_or_plain_mixed_weighted_and_plain(monkeypatch, tmp_path):
|
||||
"""Mixed list with some weighted and some unweighted values."""
|
||||
service, _ = _make_service(monkeypatch, tmp_path)
|
||||
|
||||
import random
|
||||
values = ["5::x", "y", "z"] # x has weight 5, y/z have default weight 1
|
||||
results = {"x": 0, "y": 0, "z": 0}
|
||||
for seed in range(4000):
|
||||
result = service._pick_weighted_or_plain(values, random.Random(seed))
|
||||
assert result in results
|
||||
assert "::" not in result
|
||||
results[result] += 1
|
||||
|
||||
# x (5) vs combined y+z (1+1=2) → ~71% / ~29%
|
||||
x_pct = results["x"] / sum(results.values())
|
||||
assert 0.65 < x_pct < 0.78, f"x proportion out of range: {x_pct:.3f}"
|
||||
|
||||
|
||||
def test_pick_weighted_or_plain_invalid_weight_prefix(monkeypatch, tmp_path):
|
||||
"""Invalid numeric prefix (e.g. 1.2.3) is NOT treated as a weight and
|
||||
the prefix is NOT stripped, matching the updated strict regex."""
|
||||
service, _ = _make_service(monkeypatch, tmp_path)
|
||||
|
||||
import random
|
||||
rng = random.Random(42)
|
||||
|
||||
# "1.2.3::a" is not a valid number → treated as plain text value
|
||||
result = service._pick_weighted_or_plain(["1.2.3::a", "b"], rng)
|
||||
# It should keep the full text including :: because the prefix isn't a
|
||||
# valid numeric weight according to the strict regex
|
||||
assert result == "1.2.3::a" or result == "b"
|
||||
|
||||
|
||||
def test_pick_weighted_or_plain_glob_aggregation(monkeypatch, tmp_path):
|
||||
"""Weighted wildcard resolution through glob aggregation (__*__)."""
|
||||
service, wildcards_dir = _make_service(monkeypatch, tmp_path)
|
||||
wildcards_dir.mkdir()
|
||||
|
||||
(wildcards_dir / "animals").mkdir()
|
||||
(wildcards_dir / "animals" / "cat.txt").write_text("3::tabby\n1::persian\n", encoding="utf-8")
|
||||
(wildcards_dir / "animals" / "dog.txt").write_text("retriever\npoodle\n", encoding="utf-8")
|
||||
|
||||
# __animals/*__ aggregates all values across both files
|
||||
# Weighted values should have :: stripped
|
||||
results = {"tabby": 0, "persian": 0, "retriever": 0, "poodle": 0}
|
||||
for seed in range(4000):
|
||||
expanded = service.expand_text("__animals/*__", seed=seed)
|
||||
assert expanded in results, f"Unexpected result: {expanded!r}"
|
||||
assert "::" not in expanded
|
||||
results[expanded] += 1
|
||||
|
||||
# tabby (3) vs persian (1) → ~75% / ~25% within the cat subset
|
||||
cat_total = results["tabby"] + results["persian"]
|
||||
if cat_total > 0:
|
||||
tabby_pct = results["tabby"] / cat_total
|
||||
assert 0.65 < tabby_pct < 0.85, f"tabby proportion out of range: {tabby_pct:.3f}"
|
||||
|
||||
Reference in New Issue
Block a user