Restore release-era character ref routing
This commit is contained in:
+53
-149
@@ -4454,52 +4454,32 @@ def _slot_refs_for_text(text, ref_slots):
|
||||
for slot_number in picture_tags(text):
|
||||
if not (1 <= slot_number <= len(ref_slots or [])):
|
||||
continue
|
||||
ref = ref_slots[slot_number - 1]
|
||||
raw = ref_slots[slot_number - 1]
|
||||
ref = _reference_slot(raw, slot_number)
|
||||
if _reference_image(ref) is None:
|
||||
continue
|
||||
refs.append((slot_number, ref))
|
||||
return refs
|
||||
|
||||
|
||||
def _named_refs_for_text(text, ref_slots, kinds=None):
|
||||
def _named_character_refs_for_text(text, ref_slots):
|
||||
haystack = str(text or "")
|
||||
matched = []
|
||||
wanted = {str(k).strip().lower() for k in (kinds or ()) if str(k).strip()}
|
||||
by_name = {}
|
||||
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:
|
||||
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:
|
||||
continue
|
||||
for name in _reference_name_keys(ref):
|
||||
by_name.setdefault(name.lower(), []).append((slot_number, ref, name))
|
||||
for entries in by_name.values():
|
||||
if len(entries) != 1:
|
||||
continue
|
||||
slot_number, ref, name = entries[0]
|
||||
if re.search(r"\b" + re.escape(name) + r"\b", haystack, re.I):
|
||||
matched.append((slot_number, ref))
|
||||
if re.search(r"\b" + re.escape(name) + r"\b", haystack, re.I):
|
||||
matched.append((slot_number, ref))
|
||||
break
|
||||
return matched
|
||||
|
||||
|
||||
def _matched_reference_slots(text, ref_slots, normalized_refs=None):
|
||||
def _matched_reference_slots(text, ref_slots):
|
||||
matched = []
|
||||
seen = set()
|
||||
normalized_slots = normalized_refs if normalized_refs is not None else _normalized_ref_slots(ref_slots)
|
||||
for slot_number, ref in _slot_refs_for_text(text, normalized_slots):
|
||||
if slot_number in seen:
|
||||
continue
|
||||
seen.add(slot_number)
|
||||
matched.append((slot_number, ref))
|
||||
# Character names should win before location names when a beat mentions both.
|
||||
# That keeps the face conditioning in front of the scene conditioning instead
|
||||
# of letting an early location slot dominate the named-ref list.
|
||||
for slot_number, ref in _named_refs_for_text(text, normalized_slots, kinds=("character",)):
|
||||
if slot_number in seen:
|
||||
continue
|
||||
seen.add(slot_number)
|
||||
matched.append((slot_number, ref))
|
||||
for slot_number, ref in _named_refs_for_text(text, normalized_slots, kinds=("location",)):
|
||||
for slot_number, ref in _slot_refs_for_text(text, ref_slots) + _named_character_refs_for_text(text, ref_slots):
|
||||
if slot_number in seen:
|
||||
continue
|
||||
seen.add(slot_number)
|
||||
@@ -4507,20 +4487,12 @@ def _matched_reference_slots(text, ref_slots, normalized_refs=None):
|
||||
return matched
|
||||
|
||||
|
||||
def _reference_context_for_text(
|
||||
text,
|
||||
ref_slots,
|
||||
include_character_wardrobe=True,
|
||||
normalized_refs=None,
|
||||
):
|
||||
def _reference_context_for_text(text, ref_slots):
|
||||
parts = []
|
||||
if normalized_refs is None:
|
||||
normalized_refs = _normalized_ref_slots(ref_slots)
|
||||
for slot_number, ref in _matched_reference_slots(text, ref_slots, normalized_refs=normalized_refs):
|
||||
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}"
|
||||
description = _reference_sentence(ref.get("description"))
|
||||
wardrobe = (_reference_sentence(ref.get("wardrobe"))
|
||||
if include_character_wardrobe else "")
|
||||
wardrobe = _reference_sentence(ref.get("wardrobe"))
|
||||
general = _reference_sentence(ref.get("general"))
|
||||
facts = _reference_fact_sentence(ref, label)
|
||||
if ref.get("kind") == "location":
|
||||
@@ -4557,12 +4529,10 @@ def _inject_reference_context(block, context):
|
||||
)
|
||||
|
||||
|
||||
def _reference_character_memory(ref_slots, normalized_refs=None):
|
||||
def _reference_character_memory(ref_slots):
|
||||
lines = []
|
||||
seen = set()
|
||||
if normalized_refs is None:
|
||||
normalized_refs = _normalized_ref_slots(ref_slots)
|
||||
for slot_number, ref in enumerate(normalized_refs, 1):
|
||||
for slot_number, ref in enumerate(ref_slots or [], 1):
|
||||
if ref is None:
|
||||
continue
|
||||
if ref.get("kind") != "character":
|
||||
@@ -4583,14 +4553,12 @@ def _reference_character_memory(ref_slots, normalized_refs=None):
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _connected_refs(ref_slots, normalized_refs=None):
|
||||
def _connected_refs(ref_slots):
|
||||
"""Connected refs only, preserving slot order and skipping empty sockets."""
|
||||
if normalized_refs is None:
|
||||
normalized_refs = _normalized_ref_slots(ref_slots)
|
||||
return [ref for ref in normalized_refs 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, normalized_refs=None):
|
||||
def resolve_tagged_refs(text, ref_list):
|
||||
"""(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
|
||||
@@ -4602,9 +4570,7 @@ def resolve_tagged_refs(text, ref_list, normalized_refs=None):
|
||||
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)
|
||||
if normalized_refs is 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]
|
||||
live = [n for n in wanted if 1 <= n <= len(ref_list or []) and ref_list[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)}
|
||||
|
||||
@@ -4617,10 +4583,10 @@ def resolve_tagged_refs(text, ref_list, normalized_refs=None):
|
||||
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(), [normalized_refs[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, include_named=True, normalized_refs=None):
|
||||
def resolve_prompt_refs(text, ref_list):
|
||||
"""(rewritten text, refs, dropped) for the refs a shot actually carries.
|
||||
|
||||
Explicit <Picture N> tags still decide which slot numbers the prompt points at,
|
||||
@@ -4628,70 +4594,19 @@ def resolve_prompt_refs(text, ref_list, include_named=True, normalized_refs=None
|
||||
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."""
|
||||
if normalized_refs is None:
|
||||
normalized_refs = _normalized_ref_slots(ref_list)
|
||||
rewritten, tagged_refs, dropped = resolve_tagged_refs(text, normalized_refs, normalized_refs=normalized_refs)
|
||||
rewritten, tagged_refs, dropped = resolve_tagged_refs(text, ref_list)
|
||||
refs = list(tagged_refs)
|
||||
seen = {id(ref) for ref in refs}
|
||||
if include_named:
|
||||
# Keep named character refs ahead of location refs so the identity image
|
||||
# is the first named reference the model sees on untagged beats.
|
||||
for _slot_number, ref in _named_refs_for_text(rewritten, normalized_refs, kinds=("character",)):
|
||||
marker = id(ref)
|
||||
if marker in seen:
|
||||
continue
|
||||
seen.add(marker)
|
||||
refs.append(ref)
|
||||
for _slot_number, ref in _named_refs_for_text(rewritten, normalized_refs, kinds=("location",)):
|
||||
marker = id(ref)
|
||||
if marker in seen:
|
||||
continue
|
||||
seen.add(marker)
|
||||
refs.append(ref)
|
||||
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
|
||||
|
||||
|
||||
def resolve_shot_references(text, ref_list, ref_mode, shot_index, handoff, connected_refs=None, normalized_refs=None):
|
||||
"""(rewritten text, refs, dropped, shot_tag_driven, effective_mode) for one beat.
|
||||
|
||||
Each beat should be able to pick up character names on its own. A stray
|
||||
<Picture N> tag in some other beat should not make name-matched characters
|
||||
disappear from the untagged beats in the same chain."""
|
||||
# Keep name matching local to the beat. Earlier versions used a chain-wide
|
||||
# "any tags anywhere" gate here, which accidentally turned off character
|
||||
# identity matching on untagged beats as soon as one tagged beat existed.
|
||||
# That regression is exactly what this helper prevents.
|
||||
refs = connected_refs if connected_refs is not None else _connected_refs(ref_list, normalized_refs=normalized_refs)
|
||||
if not refs:
|
||||
return text, [], [], False, ref_mode
|
||||
normalized_refs = normalized_refs if normalized_refs is not None else _normalized_ref_slots(ref_list)
|
||||
shot_tags = picture_tags(text)
|
||||
if ref_mode == "auto ref2v":
|
||||
rewritten, shot_refs, dropped = resolve_prompt_refs(
|
||||
text,
|
||||
ref_list,
|
||||
include_named=True,
|
||||
normalized_refs=normalized_refs,
|
||||
)
|
||||
if shot_tags:
|
||||
return rewritten, shot_refs, dropped, True, "auto ref2v"
|
||||
if shot_refs:
|
||||
return rewritten, shot_refs, dropped, False, "auto ref2v"
|
||||
return text, shot_references(ref_list, "every shot", shot_index, handoff, connected_refs=refs), [], False, "every shot"
|
||||
if ref_mode == "where tagged":
|
||||
if shot_tags:
|
||||
rewritten, shot_refs, dropped = resolve_prompt_refs(
|
||||
text,
|
||||
ref_list,
|
||||
include_named=False,
|
||||
normalized_refs=normalized_refs,
|
||||
)
|
||||
return rewritten, shot_refs, dropped, True, "where tagged"
|
||||
return text, shot_references(ref_list, "first shot", shot_index, handoff, connected_refs=refs), [], False, "first shot"
|
||||
return text, shot_references(ref_list, ref_mode, shot_index, handoff, connected_refs=refs), [], False, ref_mode
|
||||
|
||||
|
||||
def shot_references(ref_list, ref_mode, shot_index, handoff, connected_refs=None):
|
||||
def shot_references(ref_list, ref_mode, shot_index, handoff):
|
||||
"""Pure: which reference images shot `shot_index` is conditioned on, or [] when
|
||||
the shot should use the keyframe handoff instead.
|
||||
|
||||
@@ -4717,7 +4632,7 @@ def shot_references(ref_list, ref_mode, shot_index, handoff, connected_refs=None
|
||||
back as a soft signal (the model is shown where the last
|
||||
shot ended rather than told to start exactly there), and it
|
||||
stays a single ref2va task, so nothing conflicts."""
|
||||
refs = connected_refs if connected_refs is not None else _connected_refs(ref_list)
|
||||
refs = _connected_refs(ref_list)
|
||||
if not refs:
|
||||
return []
|
||||
if ref_mode == "auto ref2v":
|
||||
@@ -6766,8 +6681,6 @@ class H3LongVideos:
|
||||
context = _reference_context_for_text(
|
||||
block,
|
||||
ref_slots,
|
||||
include_character_wardrobe=not bool(explicit_character_memory),
|
||||
normalized_refs=normalized_ref_slots,
|
||||
)
|
||||
if context:
|
||||
block = _inject_reference_context(block, context)
|
||||
@@ -6856,20 +6769,16 @@ class H3LongVideos:
|
||||
effective_modes = []
|
||||
for shot_index, gen in enumerate(gens):
|
||||
shot_mode = beat_ref_mode_directive(beats[shot_index] if shot_index < len(beats) else "") or ref_mode
|
||||
_rewritten, refs, _dropped, shot_tag_driven, mode_eff = resolve_shot_references(
|
||||
gen,
|
||||
ref_slots,
|
||||
shot_mode,
|
||||
shot_index,
|
||||
1 if shot_index else None,
|
||||
connected_refs=connected_refs,
|
||||
normalized_refs=normalized_ref_slots,
|
||||
)
|
||||
effective_modes.append(mode_eff)
|
||||
if shot_tag_driven:
|
||||
tagged_used = tagged_used or bool(refs)
|
||||
if refs:
|
||||
on.append(shot_index + 1)
|
||||
if shot_mode in ("where tagged", "auto ref2v") and any_tags_anywhere:
|
||||
if resolve_prompt_refs(gen, ref_slots)[1]:
|
||||
on.append(shot_index + 1)
|
||||
tagged_used = True
|
||||
else:
|
||||
mode_eff = ("every shot" if shot_mode == "auto ref2v"
|
||||
else "first shot" if shot_mode == "where tagged" else shot_mode)
|
||||
effective_modes.append(mode_eff)
|
||||
if shot_references(ref_slots, mode_eff, shot_index, 1 if shot_index else None):
|
||||
on.append(shot_index + 1)
|
||||
if tagged_used:
|
||||
how = "placed by <Picture N> tags"
|
||||
else:
|
||||
@@ -6928,7 +6837,6 @@ class H3LongVideos:
|
||||
ref_mode_used = []
|
||||
continuity_used = []
|
||||
ref_aug_used = []
|
||||
any_explicit_picture_tags = False
|
||||
shot_timings = []
|
||||
if cleanup_between_shots:
|
||||
_deep_cleanup() # start the first (heaviest) shot with max free VRAM
|
||||
@@ -6946,27 +6854,23 @@ class H3LongVideos:
|
||||
shot_continuity = beat_continuity_directive(beat_text) or "auto"
|
||||
shot_ref_noise_aug = beat_ref_noise_aug_directive(beat_text)
|
||||
shot_aug = ref_noise_aug if shot_ref_noise_aug is None else shot_ref_noise_aug
|
||||
gen_prompt, shot_refs, dropped, shot_tag_driven, shot_mode_eff = resolve_shot_references(
|
||||
gen_prompt,
|
||||
ref_slots,
|
||||
shot_mode,
|
||||
i,
|
||||
handoff,
|
||||
connected_refs=connected_refs,
|
||||
normalized_refs=normalized_ref_slots,
|
||||
)
|
||||
shot_tag_driven = bool(connected_ref_count) and shot_mode in ("where tagged", "auto ref2v") and any_tags_anywhere
|
||||
if shot_tag_driven:
|
||||
shot_mode_eff = shot_mode
|
||||
else:
|
||||
shot_mode_eff = ("every shot" if shot_mode == "auto ref2v"
|
||||
else "first shot" if shot_mode == "where tagged" else shot_mode)
|
||||
carry_keyframe = False # tagged shot keeps its handoff as a keyframe
|
||||
if shot_tag_driven:
|
||||
any_explicit_picture_tags = True
|
||||
# The prompt itself says where each reference belongs: the shot whose
|
||||
# text names <Picture N> gets image N, renumbered to match what that
|
||||
# shot actually carries. Named character matches are already included
|
||||
# in the same pass; the tag only makes the slot mapping explicit.
|
||||
# Do not reintroduce a chain-wide tag gate here. Named character
|
||||
# refs must still route on beats that do not use <Picture N>.
|
||||
# shot actually carries. Every untagged shot keeps its handoff.
|
||||
gen_prompt, shot_refs, dropped = resolve_prompt_refs(gen_prompt, ref_slots)
|
||||
for n in dropped:
|
||||
if n not in ref_missing:
|
||||
ref_missing.append(n)
|
||||
else:
|
||||
shot_refs = shot_references(ref_slots, shot_mode_eff, i, handoff)
|
||||
# A shot that follows a strip starts FRESH. Continuing from a frame that
|
||||
# still shows the garment is how it reappears -- the picture outvotes the
|
||||
# text every time. Costs a cut exactly where the state changes, which is
|
||||
@@ -7026,7 +6930,7 @@ class H3LongVideos:
|
||||
"soft carry" if shot_handoff is not None else "hard cut")
|
||||
if carry_keyframe and (i + 1) not in ref_keyframed and handoff is not None:
|
||||
ref_keyframed.append(i + 1)
|
||||
ref_mode_used.append(shot_mode_eff)
|
||||
ref_mode_used.append(shot_mode if shot_tag_driven else shot_mode_eff)
|
||||
continuity_used.append(continuity_label)
|
||||
ref_aug_used.append(shot_aug)
|
||||
if shot_refs:
|
||||
@@ -7299,7 +7203,7 @@ class H3LongVideos:
|
||||
if connected_ref_count and ref_shots:
|
||||
kept = [n for n in range(1, len(gens) + 1) if n not in ref_shots]
|
||||
distinct_ref_modes = list(dict.fromkeys(ref_mode_used))
|
||||
tagged_used = any_explicit_picture_tags
|
||||
tagged_used = any(mode in ("where tagged", "auto ref2v") for mode in ref_mode_used) and any_tags_anywhere
|
||||
ref_placement = ("placed by <Picture N> tags" if tagged_used else
|
||||
f"ref_mode '{distinct_ref_modes[0]}'" if len(distinct_ref_modes) == 1 else
|
||||
"mixed per-shot ref_mode")
|
||||
|
||||
@@ -458,42 +458,6 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
|
||||
self.assertEqual([self.module._reference_image(ref) for ref in references], ["img2"])
|
||||
self.assertEqual(dropped, [])
|
||||
|
||||
def test_resolve_shot_references_uses_named_characters_without_picture_tags(self):
|
||||
refs = [
|
||||
{"kind": "character", "image": "img1", "name": "Mara"},
|
||||
{"kind": "character", "image": "img2", "name": "Jon"},
|
||||
{"kind": "location", "image": "img3", "name": "Hangar"},
|
||||
]
|
||||
|
||||
text, references, dropped, shot_tag_driven, mode_eff = self.module.resolve_shot_references(
|
||||
"[Generation 1] Mara crosses the hangar.",
|
||||
refs,
|
||||
"auto ref2v",
|
||||
0,
|
||||
None,
|
||||
)
|
||||
|
||||
self.assertEqual(text, "[Generation 1] Mara crosses the hangar.")
|
||||
self.assertEqual([self.module._reference_image(ref) for ref in references], ["img1"])
|
||||
self.assertEqual(dropped, [])
|
||||
self.assertFalse(shot_tag_driven)
|
||||
self.assertEqual(mode_eff, "auto ref2v")
|
||||
|
||||
def test_resolve_prompt_refs_prioritizes_characters_before_locations(self):
|
||||
refs = [
|
||||
{"kind": "location", "image": "img1", "name": "Hangar"},
|
||||
{"kind": "character", "image": "img2", "name": "Mara"},
|
||||
]
|
||||
|
||||
text, references, dropped = self.module.resolve_prompt_refs(
|
||||
"[Generation 1] Mara waits in the hangar.",
|
||||
refs,
|
||||
)
|
||||
|
||||
self.assertEqual(text, "[Generation 1] Mara waits in the hangar.")
|
||||
self.assertEqual([self.module._reference_image(ref) for ref in references], ["img2", "img1"])
|
||||
self.assertEqual(dropped, [])
|
||||
|
||||
def test_shot_references_uses_all_connected_sparse_slots(self):
|
||||
refs = [
|
||||
None,
|
||||
|
||||
Reference in New Issue
Block a user