Files
JasonHoku-ComfyUI-Ultimate-…/tests/test_regional_encoding.py
T
Jason HokuandClaude Fable 5.1 0e563a19cb feat(prompts): COUPLE regional prompts via native masked conditioning (#11)
`base COUPLE left | right` (2-4 equal columns; COUPLE(rows) for bands; region ` :w` weight)
splits a prompt into an unmasked base cond plus one masked cond per region, blended by
ComfyUI's own mask/mask_strength path (latent-couple style, one extra UNet pass per region,
works for every model family). COUPLE is parsed first; each part still takes AND/CAT/AVG/BREAK.

- regional_prompt.py: parser, join helpers (byte-identical output without COUPLE, cache keys
  depend on it), size-independent region tags, latent-size mask builder with tile/crop frames,
  apply_regional_masks at sample time
- batch_encoding: COUPLE prompts take the combinator path (never the batched TE path)
- conditioning_cache: conditioning_v2 keeps every entry and extra key (the old format silently
  dropped AND prompts); legacy entries are a one-time miss
- trigger words, model suffix and HiRes append_end land in the base, not the last region
- tiled HiRes and Florence2 crops pass their frame; invalid syntax stops the run early;
  LTX video warns; Builder tooltip + README section + ProjectStructure/AI-CONTEXT
- 107 new tests (parser, masks, encoding, cache round trip, source guards)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
2026-11-04 17:50:15 -08:00

125 lines
4.7 KiB
Python

"""COUPLE prompts through the real batch_encoding.encode_prompt_with_combinators.
Fake CLIP: tokenize() returns the text, encode_from_tokens() returns {"cond": text, "pooled_output": None},
so every conditioning entry is traceable by its text. CAT / AVG need real tensors (requires_torch).
"""
import importlib.util
import os
import sys
import pytest
_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
_PKG = os.path.basename(_ROOT)
sys.path.insert(0, _ROOT)
_mm = sys.modules.get("comfy.model_management")
if _mm is not None and "comfy" in sys.modules and not hasattr(sys.modules["comfy"], "model_management"):
sys.modules["comfy"].model_management = _mm
_spec = importlib.util.spec_from_file_location(f"{_PKG}.batch_encoding_real_regional",
os.path.join(_ROOT, "batch_encoding.py"))
be = importlib.util.module_from_spec(_spec)
be.__package__ = _PKG
_spec.loader.exec_module(be)
import importlib # noqa: E402
rp = importlib.import_module(f"{_PKG}.regional_prompt")
class FakeCLIP:
def __init__(self):
self.tokenized = []
def tokenize(self, text):
self.tokenized.append(text)
return text
def encode_from_tokens(self, tokens, return_pooled=True, return_dict=True):
return {"cond": tokens, "pooled_output": None}
def _tags(out):
return [e[1].get(rp.REGION_KEY) for e in out]
def test_base_and_two_regions_in_order_with_tags():
out = be.encode_prompt_with_combinators(FakeCLIP(), "base COUPLE a | b")
assert [e[0] for e in out] == ["base", "a", "b"]
assert _tags(out) == [None, rp.region_tag(0, 2, "cols", 1.0), rp.region_tag(1, 2, "cols", 1.0)]
def test_rows_layout_and_weight_in_tag():
out = be.encode_prompt_with_combinators(FakeCLIP(), "base COUPLE(rows) top | bottom :1.5")
assert _tags(out)[1:] == [rp.region_tag(0, 2, "rows", 1.0), rp.region_tag(1, 2, "rows", 1.5)]
def test_and_inside_region_gives_entries_with_same_tag():
out = be.encode_prompt_with_combinators(FakeCLIP(), "base COUPLE a AND a2 | b")
assert [e[0] for e in out] == ["base", "a", "a2", "b"]
assert _tags(out)[1] == _tags(out)[2] == rp.region_tag(0, 2, "cols", 1.0)
assert _tags(out)[3] == rp.region_tag(1, 2, "cols", 1.0)
def test_empty_base_has_no_base_entry():
out = be.encode_prompt_with_combinators(FakeCLIP(), "COUPLE a | b")
assert [e[0] for e in out] == ["a", "b"]
assert all(t is not None for t in _tags(out))
def test_empty_and_zero_regions_keep_slot_but_add_no_entry():
out = be.encode_prompt_with_combinators(FakeCLIP(), "base COUPLE a | | c :0 | d")
assert [e[0] for e in out] == ["base", "a", "d"]
assert [t["index"] for t in _tags(out)[1:]] == [0, 3]
assert all(t["count"] == 4 for t in _tags(out)[1:])
def test_base_with_break_goes_through_tokenizer_untouched():
clip = FakeCLIP()
out = be.encode_prompt_with_combinators(clip, "x BREAK y COUPLE a | b")
assert out[0][0] == "x BREAK y"
def test_tags_are_independent_dicts_and_pooled_kept():
out = be.encode_prompt_with_combinators(FakeCLIP(), "COUPLE a AND b | c")
assert out[0][1][rp.REGION_KEY] is not out[1][1][rp.REGION_KEY]
assert "pooled_output" in out[0][1]
def test_has_combinators_true_for_couple():
assert be._has_combinators("x COUPLE a | b")
assert not be._has_combinators("x couple a | b")
assert not be._has_combinators("plain prompt")
def test_batchable_ids_hands_couple_back_as_combinator():
assert be._batchable_ids(object(), "x COUPLE a | b", "clip_l", 0) == (None, "combinator")
@pytest.mark.parametrize("p", ["x COUPLE only one", "x COUPLE(diag) a | b", "x COUPLE a | b COUPLE c | d"])
def test_invalid_prompt_raises(p):
with pytest.raises(rp.RegionalPromptError):
be.encode_prompt_with_combinators(FakeCLIP(), p)
def test_plain_prompt_still_uses_single_encode():
clip = FakeCLIP()
out = be.encode_prompt_with_combinators(clip, "1girl, couple, solo")
assert [e[0] for e in out] == ["1girl, couple, solo"]
assert rp.REGION_KEY not in out[0][1]
@pytest.mark.requires_torch
def test_cat_and_avg_inside_regions():
import torch
class TClip(FakeCLIP):
def encode_from_tokens(self, tokens, return_pooled=True, return_dict=True):
return {"cond": torch.full((1, 2, 3), float(len(tokens))), "pooled_output": None}
out = be.encode_prompt_with_combinators(TClip(), "base COUPLE aa CAT bbb | c AVG(0.5) dddd")
assert len(out) == 3
assert out[0][1].get(rp.REGION_KEY) is None
assert out[1][1][rp.REGION_KEY]["index"] == 0 and out[2][1][rp.REGION_KEY]["index"] == 1
assert out[1][0].shape[1] == 4 # CAT concatenates along the token axis
assert torch.allclose(out[2][0], torch.full((1, 2, 3), 2.5)) # AVG 0.5 of len 1 and len 4