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
+123 -70
View File
@@ -4462,54 +4462,82 @@ def _slot_refs_for_text(text, ref_slots):
return refs return refs
def _named_character_refs_for_text(text, ref_slots): def _named_character_refs_for_text(text, ref_slots):
haystack = str(text or "") return _named_refs_for_text(text, ref_slots, kinds=("character",))
matched = []
for slot_number, raw in enumerate(ref_slots or [], 1):
ref = _reference_slot(raw, slot_number) def _named_refs_for_text(text, ref_slots, kinds=("character",)):
if ref.get("kind") != "character" or _reference_image(ref) is None: """Backward-compatible alias for older call sites.
continue
for name in _reference_name_keys(ref): The script annotation path still expects the shorter helper name, while the
if re.search(r"\b" + re.escape(name) + r"\b", haystack, re.I): newer reference matcher keeps the more explicit character-specific helper.
matched.append((slot_number, ref)) """
break kinds = tuple(dict.fromkeys(kinds or ()))
return matched if not kinds:
return []
haystack = str(text or "")
kind_hits = {kind: [] for kind in kinds}
name_hits = {kind: {} for kind in kinds}
for slot_number, raw in enumerate(ref_slots or [], 1):
ref = _reference_slot(raw, slot_number)
kind = ref.get("kind")
if kind not in kind_hits or _reference_image(ref) is None:
continue
matched_names = []
for name in _reference_name_keys(ref):
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))
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)
matched.append((slot_number, ref)) matched.append((slot_number, ref))
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}"
description = _reference_sentence(ref.get("description")) description = _reference_sentence(ref.get("description"))
wardrobe = _reference_sentence(ref.get("wardrobe")) wardrobe = _reference_sentence(ref.get("wardrobe"))
general = _reference_sentence(ref.get("general")) general = _reference_sentence(ref.get("general"))
facts = _reference_fact_sentence(ref, label) facts = _reference_fact_sentence(ref, label)
if ref.get("kind") == "location": if ref.get("kind") == "location":
if description: if description:
parts.append(f"Location context for {label}: {description}") parts.append(f"Location context for {label}: {description}")
if general: if general:
parts.append(f"Location notes for {label}: {general}") parts.append(f"Location notes for {label}: {general}")
continue continue
if facts: if facts:
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}")
return " ".join(parts).strip() return " ".join(parts).strip()
def _inject_reference_context(block, context): def _inject_reference_context(block, context):
@@ -4558,8 +4586,8 @@ def _connected_refs(ref_slots):
return [ref for ref in (ref_slots or []) if _reference_image(ref) is not None] return [ref for ref in (ref_slots or []) if _reference_image(ref) is not None]
def resolve_tagged_refs(text, ref_list): def resolve_tagged_refs(text, ref_list):
"""(rewritten text, images, dropped) for the <Picture N> tags in ONE shot. """(rewritten text, images, dropped) for the <Picture N> tags in ONE shot.
The tokenizer numbers references by their position in the list it is handed, so The tokenizer numbers references by their position in the list it is handed, so
a shot that uses only <Picture 2> would receive that image labelled a shot that uses only <Picture 2> would receive that image labelled
@@ -4567,12 +4595,12 @@ def resolve_tagged_refs(text, ref_list):
RENUMBERED per shot to match what that shot actually carries: slot 2 alone RENUMBERED per shot to match what that shot actually carries: slot 2 alone
becomes <Picture 1>, slots 2 and 4 become <Picture 1> and <Picture 2>. becomes <Picture 1>, slots 2 and 4 become <Picture 1> and <Picture 2>.
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)}
def sub(m): def sub(m):
n = int(m.group(1)) n = int(m.group(1))
@@ -4586,24 +4614,49 @@ 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,
but named character matches must ride into the real ref-image list too. Without but named character matches must ride into the real ref-image list too. Without
that split, a shot could inherit the facts/context for "Mara" and "Jon" while that split, a shot could inherit the facts/context for "Mara" and "Jon" while
only carrying a tagged location image, which reads exactly like the names were only carrying a tagged location image, which reads exactly like the names were
understood but the faces were ignored.""" understood but the faces were ignored."""
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:
marker = id(ref) for _slot_number, ref in _named_refs_for_text(rewritten, ref_list, kinds=("character", "location")):
if marker in seen: marker = id(ref)
continue if marker in seen:
seen.add(marker) continue
refs.append(ref) seen.add(marker)
return rewritten, refs, dropped refs.append(ref)
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):
+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"]