Harden and streamline H3 reference matching

dumas_h3_longvideos.py

@@ -3288,9 +3288,10 @@ def annotate_script_debug(gens, anatomy_shots, anatomy_mode, ref_slots):
 def annotate_script_refs(gens, ref_slots):
     """Per-shot reference routing summary for the script socket."""
     lines = []
+    normalized_slots = _normalized_ref_slots(ref_slots)
     for shot_index, gen in enumerate(gens or [], 1):
-        tagged = _slot_refs_for_text(gen, ref_slots)
-        named = _named_character_refs_for_text(gen, ref_slots)
+        tagged = _slot_refs_for_text(gen, normalized_slots)
+        named = _named_refs_for_text(gen, normalized_slots, kinds=("character", "location"))
         seen = set()
         merged = []
         for slot_number, ref in tagged:
@@ -4282,7 +4283,10 @@ def _reference_slot(ref, slot_index=None):
 
 def _reference_image(ref):
     try:
-        normalized = _reference_slot(ref)
+        if isinstance(ref, dict):
+            normalized = _image_nodes.normalize_reference(ref, allow_image_fallback=False)
+        else:
+            normalized = _reference_slot(ref)
     except Exception:
         return None
     return normalized.get("image")
@@ -4372,25 +4376,39 @@ def _reference_name_keys(ref):
     return out
 
 
+def _normalized_ref_slots(ref_slots):
+    """Normalize every connected ref once so downstream helpers can reuse them."""
+    out = []
+    for slot_number, raw in enumerate(ref_slots or [], 1):
+        if raw is None:
+            out.append(None)
+        elif isinstance(raw, dict) and "image" in raw and raw.get("image") is None:
+            out.append(_image_nodes.normalize_reference(raw, picture_id=slot_number, allow_image_fallback=False))
+        else:
+            out.append(_reference_slot(raw, slot_number))
+    return tuple(out)
+
+
 def _slot_refs_for_text(text, ref_slots):
     refs = []
     for slot_number in picture_tags(text):
         if not (1 <= slot_number <= len(ref_slots or [])):
             continue
-        raw = ref_slots[slot_number - 1]
-        ref = _reference_slot(raw, slot_number)
+        ref = ref_slots[slot_number - 1]
         if _reference_image(ref) is None:
             continue
         refs.append((slot_number, ref))
     return refs
 
 
-def _named_character_refs_for_text(text, ref_slots):
+def _named_refs_for_text(text, ref_slots, kinds=None):
     haystack = str(text or "")
     matched = []
-    for slot_number, raw in enumerate(ref_slots or [], 1):
-        ref = _reference_slot(raw, slot_number)
-        if ref.get("kind") != "character" or _reference_image(ref) is None:
+    wanted = {str(k).strip().lower() for k in (kinds or ()) if str(k).strip()}
+    for slot_number, ref in enumerate(ref_slots or [], 1):
+        if ref is None or _reference_image(ref) is None:
+            continue
+        if wanted and str(ref.get("kind") or "").strip().lower() not in wanted:
             continue
         for name in _reference_name_keys(ref):
             if re.search(r"\b" + re.escape(name) + r"\b", haystack, re.I):
@@ -4402,7 +4420,11 @@ def _named_character_refs_for_text(text, ref_slots):
 def _matched_reference_slots(text, ref_slots):
     matched = []
     seen = set()
-    for slot_number, ref in _slot_refs_for_text(text, ref_slots) + _named_character_refs_for_text(text, ref_slots):
+    normalized_slots = _normalized_ref_slots(ref_slots)
+    for slot_number, ref in (
+        _slot_refs_for_text(text, normalized_slots)
+        + _named_refs_for_text(text, normalized_slots, kinds=("character", "location"))
+    ):
         if slot_number in seen:
             continue
         seen.add(slot_number)
@@ -4456,8 +4478,9 @@ def _inject_reference_context(block, context):
 def _reference_character_memory(ref_slots):
     lines = []
     seen = set()
-    for slot_number, raw in enumerate(ref_slots or [], 1):
-        ref = _reference_slot(raw, slot_number)
+    for slot_number, ref in enumerate(_normalized_ref_slots(ref_slots), 1):
+        if ref is None:
+            continue
         if ref.get("kind") != "character":
             continue
         wardrobe = _reference_text(ref.get("wardrobe"))
@@ -4478,7 +4501,7 @@ def _reference_character_memory(ref_slots):
 
 def _connected_refs(ref_slots):
     """Connected refs only, preserving slot order and skipping empty sockets."""
-    return [ref for ref in (ref_slots or []) if _reference_image(ref) is not None]
+    return [ref for ref in _normalized_ref_slots(ref_slots) if _reference_image(ref) is not None]
 
 
 def resolve_tagged_refs(text, ref_list):
@@ -4493,7 +4516,8 @@ def resolve_tagged_refs(text, ref_list):
     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."""
     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]
+    normalized_refs = _normalized_ref_slots(ref_list)
+    live = [n for n in wanted if 1 <= n <= len(normalized_refs) and _reference_image(normalized_refs[n - 1]) is not None]
     dropped = [n for n in wanted if n not in live]
     renumber = {old: new for new, old in enumerate(live, 1)}
 
@@ -4501,12 +4525,12 @@ def resolve_tagged_refs(text, ref_list):
         n = int(m.group(1))
         return f"<Picture {renumber[n]}>" if n in renumber else ""
 
-    out = _PICTURE_TAG.sub(sub, text or "")
+    out = _PICTURE_TAG.sub(sub, text or "")
     if dropped:                       # tidy the gap a removed tag leaves behind
         out = re.sub(r"\s+([,.;:])", r"\1", out)
         out = re.sub(r"(,\s*){2,}", ", ", out)
         out = re.sub(r"\s{2,}", " ", out)
-    return out.strip(), [ref_list[n - 1] for n in live], dropped
+    return out.strip(), [normalized_refs[n - 1] for n in live], dropped
 
 
 def resolve_prompt_refs(text, ref_list):
@@ -4517,10 +4541,11 @@ def resolve_prompt_refs(text, ref_list):
     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
     understood but the faces were ignored."""
-    rewritten, tagged_refs, dropped = resolve_tagged_refs(text, ref_list)
+    normalized_refs = _normalized_ref_slots(ref_list)
+    rewritten, tagged_refs, dropped = resolve_tagged_refs(text, normalized_refs)
     refs = list(tagged_refs)
     seen = {id(ref) for ref in refs}
-    for _slot_number, ref in _named_character_refs_for_text(rewritten, ref_list):
+    for _slot_number, ref in _named_refs_for_text(rewritten, normalized_refs, kinds=("character", "location")):
         marker = id(ref)
         if marker in seen:
             continue