Align H3 prompt refs with beat numbering

2026-09-02 15:06:41 +00:00
parent 76b88b0c59
commit ecd8ca6336
2 changed files with 137 additions and 65 deletions
@@ -4519,10 +4519,34 @@ def _matched_reference_slots(text, ref_slots):
    return matched


-def _reference_context_for_text(text, ref_slots, include_character_wardrobe=True):
+def _resolved_picture_number(ref, resolved_refs):
+    target_image = _reference_image(ref)
+    for picture_number, active_ref in enumerate(resolved_refs or [], 1):
+        if active_ref is ref:
+            return picture_number
+        active_image = _reference_image(active_ref)
+        if target_image is not None and active_image is target_image:
+            return picture_number
+    return None
+
+
+def _resolved_ref_label(ref, slot_number, resolved_refs=None):
+    name = _reference_text(ref.get("name")) or _reference_text(ref.get("id")) or f"reference {slot_number}"
+    if not resolved_refs:
+        return name
+    picture_number = _resolved_picture_number(ref, resolved_refs)
+    if picture_number is not None:
+        return f"<Picture {picture_number}> {name}"
+    return name
+
+
+def _reference_context_for_text(text, ref_slots, include_character_wardrobe=True, resolved_refs=None):
    parts = []
-    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}"
+    active_slots = tuple(resolved_refs) if resolved_refs else ref_slots
+    for slot_number, ref in _matched_reference_slots(text, active_slots):
+        if resolved_refs and _resolved_picture_number(ref, resolved_refs) is None:
+            continue
+        label = _resolved_ref_label(ref, slot_number, resolved_refs)
        description = _reference_sentence(ref.get("description"))
        wardrobe = _reference_sentence(ref.get("wardrobe"))
        general = _reference_sentence(ref.get("general"))
@@ -4664,6 +4688,18 @@ def resolve_tag_driven_prompt_refs(text, ref_list):
    return rewritten, refs, dropped


+def _resolve_beat_prompt_refs(text, ref_list, beat_text, connected_ref_count, ref_mode,
+                              any_tags_anywhere, shot_index, handoff):
+    shot_mode = beat_ref_mode_directive(beat_text) or ref_mode
+    shot_tag_driven = bool(connected_ref_count) and shot_mode in ("where tagged", "auto ref2v") and any_tags_anywhere
+    if shot_tag_driven:
+        rewritten, refs, dropped = resolve_tag_driven_prompt_refs(text, ref_list)
+        return rewritten, refs, dropped, shot_mode, True
+    shot_mode_eff = ("every shot" if shot_mode == "auto ref2v"
+                     else "first shot" if shot_mode == "where tagged" else shot_mode)
+    return text, shot_references(ref_list, shot_mode_eff, shot_index, handoff), [], shot_mode_eff, False
+
+
 def resolve_shot_references(text, ref_list, ref_mode="auto ref2v", shot_index=0, handoff=None):
    """Compatibility wrapper for older tests and helper code.

@@ -6515,6 +6551,7 @@ class H3LongVideos:
        ))
        normalized_ref_slots = _normalized_ref_slots(ref_slots)
        connected_refs = [ref for ref in normalized_ref_slots if _reference_image(ref) is not None]
+        connected_ref_count = len(connected_refs)
        direct_ref_count = len(connected_refs)
        explicit_character_memory = (character_memory or "").strip()
        derived_character_memory = _reference_character_memory(ref_slots)
@@ -6719,28 +6756,41 @@ class H3LongVideos:
                  "'seconds:' on the beat")
        wardrobe_notes = []
        strip_shots = []      # shots that newly bared a zone -> the NEXT shot starts fresh
-        gens = distribute_generations(anchor, beats, global_soundscape.strip(),
-                                      non_diegetic_music.strip(), effective_character_memory,
-                                       auto_wardrobe, auto_silence_nonspeech, allow_nonspeech_vocals, count_subjects,
-                                      lora_on, notes_out=wardrobe_notes, auto_props=auto_props,
-                                      prevent_nudity=prevent_nudity,
-                                      exposed_terms=exposed_terms, strip_out=strip_shots,
-                                      anatomy_guard=anatomy_on,
-                                      anatomy_auto=anatomy_auto,
-                                      lock_restraints=lock_restraints,
-                                      solidity_guard=solidity_guard,
-                                      motion_guard=motion_guard,
-                                      contact_guard=contact_guard,
-                                      count_auto=(subject_count_guard == "auto"))
+        raw_gens = distribute_generations(anchor, beats, global_soundscape.strip(),
+                                          non_diegetic_music.strip(), effective_character_memory,
+                                           auto_wardrobe, auto_silence_nonspeech, allow_nonspeech_vocals, count_subjects,
+                                          lora_on, notes_out=wardrobe_notes, auto_props=auto_props,
+                                          prevent_nudity=prevent_nudity,
+                                          exposed_terms=exposed_terms, strip_out=strip_shots,
+                                          anatomy_guard=anatomy_on,
+                                          anatomy_auto=anatomy_auto,
+                                          lock_restraints=lock_restraints,
+                                          solidity_guard=solidity_guard,
+                                          motion_guard=motion_guard,
+                                          contact_guard=contact_guard,
+                                          count_auto=(subject_count_guard == "auto"))
+        any_tags_anywhere = any(picture_tags(g) for g in raw_gens)
        enriched_gens = []
-        for block in gens:
-            context = _reference_context_for_text(
+        for i, block in enumerate(raw_gens):
+            beat_text = beats[i] if i < len(beats) else ""
+            rewritten_block, shot_refs, _dropped, _shot_mode_eff, _shot_tag_driven = _resolve_beat_prompt_refs(
                block,
                ref_slots,
+                beat_text,
+                connected_ref_count,
+                ref_mode,
+                any_tags_anywhere,
+                i,
+                None,
+            )
+            context = _reference_context_for_text(
+                rewritten_block,
+                ref_slots,
+                resolved_refs=shot_refs,
            )
            if context:
-                block = _inject_reference_context(block, context)
-            enriched_gens.append(block)
+                rewritten_block = _inject_reference_context(rewritten_block, context)
+            enriched_gens.append(rewritten_block)
        gens = enriched_gens

        # A scenery beat mid-chain hands the next shot a frame with no people in
@@ -6770,7 +6820,6 @@ class H3LongVideos:
                                for lbl, txt in preflight if txt)
        override_notes = [beat_override_summary(beat, index) for index, beat in enumerate(beats, 1)]
        override_notes = [note for note in override_notes if note]
-        any_tags_anywhere = any(picture_tags(g) for g in gens)
        anatomy_shots = [
            shot_index + 1
            for shot_index, gen in enumerate(gens)
@@ -6823,7 +6872,7 @@ class H3LongVideos:
                tagged_used = False
                on = []
                effective_modes = []
-                for shot_index, gen in enumerate(gens):
+                for shot_index, gen in enumerate(raw_gens):
                    shot_mode = beat_ref_mode_directive(beats[shot_index] if shot_index < len(beats) else "") or ref_mode
                    if shot_mode in ("where tagged", "auto ref2v") and any_tags_anywhere:
                        if resolve_tag_driven_prompt_refs(gen, ref_slots)[1]:
@@ -6885,7 +6934,6 @@ class H3LongVideos:
        latent_chunks = []                 # per-shot sampled latents, pre-decode
        mouth_settled = []                 # shots seeded from a settled (closed) mouth
        handoff, sr = first_frame, None
-        connected_ref_count = len(connected_refs)
        ref_shots = []                     # which shots ended up ref-conditioned
        ref_missing = []                   # <Picture N> tags naming an unconnected slot
        ref_carried = []                   # tagged shots that kept continuity as an extra ref
@@ -6898,7 +6946,7 @@ class H3LongVideos:
            _deep_cleanup()          # start the first (heaviest) shot with max free VRAM

        shot_lens = (lens + [ln] * len(gens))[:len(gens)]
-        for i, gen_prompt in enumerate(gens):
+        for i, base_prompt in enumerate(raw_gens):
            # denoise is fixed at 1.0 (partial denoise desyncs the joint AV schedule).
            sa = (seed + i if vary_seed_per_shot else seed, steps, cfg, sampler_name, scheduler, 1.0)
            ln_i = shot_lens[i]        # this beat's own length (<= the VRAM ceiling)
@@ -6906,27 +6954,30 @@ class H3LongVideos:
            # _build_shot_conditioning for how they are packed. On ComfyUI 0.31+ a
            # shot may carry BOTH references and a keyframe.
            beat_text = beats[i] if i < len(beats) else ""
-            shot_mode = beat_ref_mode_directive(beat_text) or ref_mode
            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
-            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)
+            gen_prompt, shot_refs, dropped, shot_mode_eff, shot_tag_driven = _resolve_beat_prompt_refs(
+                base_prompt,
+                ref_slots,
+                beat_text,
+                connected_ref_count,
+                ref_mode,
+                any_tags_anywhere,
+                i,
+                handoff,
+            )
            carry_keyframe = False       # tagged shot keeps its handoff as a keyframe
-            if shot_tag_driven:
-                # 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. Every untagged shot keeps its handoff.
-                gen_prompt, shot_refs, dropped = resolve_tag_driven_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)
+            for n in dropped:
+                if n not in ref_missing:
+                    ref_missing.append(n)
+            context = _reference_context_for_text(
+                gen_prompt,
+                ref_slots,
+                resolved_refs=shot_refs,
+            )
+            if context:
+                gen_prompt = _inject_reference_context(gen_prompt, context)
            # 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
@@ -714,6 +714,27 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
        self.assertIn("Location context for Hangar: wet concrete floor.", context)
        self.assertIn("Location notes for Hangar: cold industrial lighting.", context)

+    def test_reference_context_uses_resolved_picture_numbers_for_per_beat_refs(self):
+        refs = [
+            {"kind": "character", "image": "img1", "name": "Bill", "description": "very tall", "facts": {"age": "25"}},
+            {"kind": "location", "image": "img2", "name": "Pub", "description": "warm wood bar"},
+        ]
+
+        rewritten, resolved_refs, dropped = self.module.resolve_prompt_refs(
+            "[Generation 1] Bill leans on <Picture 2>.",
+            refs,
+        )
+        context = self.module._reference_context_for_text(
+            rewritten,
+            refs,
+            resolved_refs=resolved_refs,
+        )
+
+        self.assertEqual(dropped, [])
+        self.assertIn("Character facts for <Picture 2> Bill: 25 years old.", context)
+        self.assertIn("Persistent appearance for <Picture 2> Bill: very tall.", context)
+        self.assertIn("Location context for <Picture 1> Pub: warm wood bar.", context)
+
    def test_reference_context_skips_ambiguous_name_matches(self):
        refs = [
            {"kind": "character", "image": "img1", "name": "Alex", "description": "short dark hair"},