Fix named reference lookup for long videos

This commit is contained in:
2026-08-31 09:40:51 +00:00
parent a2dabaab7e
commit 130b268d1e
2 changed files with 137 additions and 70 deletions
+62 -9
View File
@@ -4463,23 +4463,51 @@ def _slot_refs_for_text(text, ref_slots):
def _named_character_refs_for_text(text, ref_slots): def _named_character_refs_for_text(text, ref_slots):
return _named_refs_for_text(text, ref_slots, kinds=("character",))
def _named_refs_for_text(text, ref_slots, kinds=("character",)):
"""Backward-compatible alias for older call sites.
The script annotation path still expects the shorter helper name, while the
newer reference matcher keeps the more explicit character-specific helper.
"""
kinds = tuple(dict.fromkeys(kinds or ()))
if not kinds:
return []
haystack = str(text or "") haystack = str(text or "")
matched = [] kind_hits = {kind: [] for kind in kinds}
name_hits = {kind: {} for kind in kinds}
for slot_number, raw in enumerate(ref_slots or [], 1): for slot_number, raw in enumerate(ref_slots or [], 1):
ref = _reference_slot(raw, slot_number) ref = _reference_slot(raw, slot_number)
if ref.get("kind") != "character" or _reference_image(ref) is None: kind = ref.get("kind")
if kind not in kind_hits or _reference_image(ref) is None:
continue continue
matched_names = []
for name in _reference_name_keys(ref): for name in _reference_name_keys(ref):
if re.search(r"\b" + re.escape(name) + r"\b", haystack, re.I): if re.search(r"\b" + re.escape(name) + r"\b", haystack, re.I):
matched_names.append(name.lower())
name_hits[kind].setdefault(name.lower(), []).append((slot_number, ref))
if matched_names:
kind_hits[kind].append((slot_number, ref, tuple(set(matched_names))))
matched = []
seen = set()
for kind in kinds:
for slot_number, ref, matched_names in kind_hits.get(kind, []):
if not any(len(name_hits[kind].get(name, [])) == 1 for name in matched_names):
continue
marker = id(ref)
if marker in seen:
continue
seen.add(marker)
matched.append((slot_number, ref)) matched.append((slot_number, ref))
break
return matched return matched
def _matched_reference_slots(text, ref_slots): def _matched_reference_slots(text, ref_slots):
matched = [] matched = []
seen = set() seen = set()
for slot_number, ref in _slot_refs_for_text(text, ref_slots) + _named_character_refs_for_text(text, ref_slots): for slot_number, ref in _slot_refs_for_text(text, ref_slots) + _named_refs_for_text(text, ref_slots, kinds=("character", "location")):
if slot_number in seen: if slot_number in seen:
continue continue
seen.add(slot_number) seen.add(slot_number)
@@ -4487,7 +4515,7 @@ def _matched_reference_slots(text, ref_slots):
return matched return matched
def _reference_context_for_text(text, ref_slots): def _reference_context_for_text(text, ref_slots, include_character_wardrobe=True):
parts = [] parts = []
for slot_number, ref in _matched_reference_slots(text, ref_slots): for slot_number, ref in _matched_reference_slots(text, ref_slots):
label = _reference_text(ref.get("name")) or _reference_text(ref.get("id")) or f"reference {slot_number}" label = _reference_text(ref.get("name")) or _reference_text(ref.get("id")) or f"reference {slot_number}"
@@ -4505,7 +4533,7 @@ def _reference_context_for_text(text, ref_slots):
parts.append(facts) parts.append(facts)
if description: if description:
parts.append(f"Persistent appearance for {label}: {description}") parts.append(f"Persistent appearance for {label}: {description}")
if wardrobe: if include_character_wardrobe and wardrobe:
parts.append(f"Persistent wardrobe/style for {label}: {wardrobe}") parts.append(f"Persistent wardrobe/style for {label}: {wardrobe}")
if general: if general:
parts.append(f"Character notes for {label}: {general}") parts.append(f"Character notes for {label}: {general}")
@@ -4570,7 +4598,7 @@ def resolve_tagged_refs(text, ref_list):
A tag naming a slot with no image connected refers to nothing at all, so it is A tag naming a slot with no image connected refers to nothing at all, so it is
removed from the text rather than left to confuse the encoder, and reported.""" removed from the text rather than left to confuse the encoder, and reported."""
wanted = picture_tags(text) wanted = picture_tags(text)
live = [n for n in wanted if 1 <= n <= len(ref_list or []) and ref_list[n - 1] is not None] live = [n for n in wanted if 1 <= n <= len(ref_list or []) and _reference_image(ref_list[n - 1]) is not None]
dropped = [n for n in wanted if n not in live] dropped = [n for n in wanted if n not in live]
renumber = {old: new for new, old in enumerate(live, 1)} renumber = {old: new for new, old in enumerate(live, 1)}
@@ -4586,7 +4614,7 @@ def resolve_tagged_refs(text, ref_list):
return out.strip(), [ref_list[n - 1] for n in live], dropped return out.strip(), [ref_list[n - 1] for n in live], dropped
def resolve_prompt_refs(text, ref_list): def resolve_prompt_refs(text, ref_list, include_named=True):
"""(rewritten text, refs, dropped) for the refs a shot actually carries. """(rewritten text, refs, dropped) for the refs a shot actually carries.
Explicit <Picture N> tags still decide which slot numbers the prompt points at, Explicit <Picture N> tags still decide which slot numbers the prompt points at,
@@ -4597,7 +4625,8 @@ def resolve_prompt_refs(text, ref_list):
rewritten, tagged_refs, dropped = resolve_tagged_refs(text, ref_list) rewritten, tagged_refs, dropped = resolve_tagged_refs(text, ref_list)
refs = list(tagged_refs) refs = list(tagged_refs)
seen = {id(ref) for ref in refs} seen = {id(ref) for ref in refs}
for _slot_number, ref in _named_character_refs_for_text(rewritten, ref_list): if include_named:
for _slot_number, ref in _named_refs_for_text(rewritten, ref_list, kinds=("character", "location")):
marker = id(ref) marker = id(ref)
if marker in seen: if marker in seen:
continue continue
@@ -4606,6 +4635,30 @@ def resolve_prompt_refs(text, ref_list):
return rewritten, refs, dropped return rewritten, refs, dropped
def resolve_shot_references(text, ref_list, ref_mode="auto ref2v", shot_index=0, handoff=None):
"""Compatibility wrapper for older tests and helper code.
The renderer now uses `resolve_prompt_refs` plus `shot_references` directly,
but some helper tests still check the combined resolution path."""
if ref_mode == "where tagged":
rewritten, refs, dropped = resolve_prompt_refs(text, ref_list, include_named=False)
return rewritten, refs, dropped, bool(picture_tags(text)), ref_mode
if ref_mode == "auto ref2v":
rewritten, tagged_refs, dropped = resolve_tagged_refs(text, ref_list)
refs = list(tagged_refs)
seen = {id(ref) for ref in refs}
for _slot_number, ref in _named_character_refs_for_text(rewritten, ref_list):
marker = id(ref)
if marker in seen:
continue
seen.add(marker)
refs.append(ref)
return rewritten, refs, dropped, bool(picture_tags(text)), ref_mode
rewritten, _, dropped = resolve_tagged_refs(text, ref_list)
refs = shot_references(ref_list, ref_mode, shot_index, handoff)
return rewritten, refs, dropped, False, ref_mode
def shot_references(ref_list, ref_mode, shot_index, handoff): def shot_references(ref_list, ref_mode, shot_index, handoff):
"""Pure: which reference images shot `shot_index` is conditioned on, or [] when """Pure: which reference images shot `shot_index` is conditioned on, or [] when
the shot should use the keyframe handoff instead. the shot should use the keyframe handoff instead.
+14
View File
@@ -513,6 +513,20 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
["img2", "img4", "img7"], ["img2", "img4", "img7"],
) )
def test_annotate_script_refs_handles_named_characters_and_locations(self):
refs = [
{"kind": "character", "image": "img1", "name": "Mara"},
{"kind": "location", "image": "img2", "name": "Hangar"},
]
report = self.module.annotate_script_refs(
["Mara waits in the Hangar.", "Nobody else is here."],
refs,
)
self.assertIn("# shot 1 refs: Picture 1 Mara (by name); Picture 2 Hangar (by name)", report)
self.assertIn("# shot 2 refs: none", report)
def test_input_types_expose_nine_ref_slots(self): def test_input_types_expose_nine_ref_slots(self):
optional = self.module.H3LongVideos.INPUT_TYPES()["optional"] optional = self.module.H3LongVideos.INPUT_TYPES()["optional"]