Align H3 prompt refs with beat numbering

dumas_h3_longvideos.py

@@ -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):
+    active_slots = tuple(resolved_refs) if resolved_refs else ref_slots
-        label = _reference_text(ref.get("name")) or _reference_text(ref.get("id")) or f"reference {slot_number}"
+    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(),
+        raw_gens = distribute_generations(anchor, beats, global_soundscape.strip(),
-                                      non_diegetic_music.strip(), effective_character_memory,
+                                          non_diegetic_music.strip(), effective_character_memory,
-                                       auto_wardrobe, auto_silence_nonspeech, allow_nonspeech_vocals, count_subjects,
+                                           auto_wardrobe, auto_silence_nonspeech, allow_nonspeech_vocals, count_subjects,
-                                      lora_on, notes_out=wardrobe_notes, auto_props=auto_props,
+                                          lora_on, notes_out=wardrobe_notes, auto_props=auto_props,
-                                      prevent_nudity=prevent_nudity,
+                                          prevent_nudity=prevent_nudity,
-                                      exposed_terms=exposed_terms, strip_out=strip_shots,
+                                          exposed_terms=exposed_terms, strip_out=strip_shots,
-                                      anatomy_guard=anatomy_on,
+                                          anatomy_guard=anatomy_on,
-                                      anatomy_auto=anatomy_auto,
+                                          anatomy_auto=anatomy_auto,
-                                      lock_restraints=lock_restraints,
+                                          lock_restraints=lock_restraints,
-                                      solidity_guard=solidity_guard,
+                                          solidity_guard=solidity_guard,
-                                      motion_guard=motion_guard,
+                                          motion_guard=motion_guard,
-                                      contact_guard=contact_guard,
+                                          contact_guard=contact_guard,
-                                      count_auto=(subject_count_guard == "auto"))
+                                          count_auto=(subject_count_guard == "auto"))
+        any_tags_anywhere = any(picture_tags(g) for g in raw_gens)
         enriched_gens = []
-        for block in gens:
+        for i, block in enumerate(raw_gens):
-            context = _reference_context_for_text(
+            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)
+                rewritten_block = _inject_reference_context(rewritten_block, context)
-            enriched_gens.append(block)
+            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
+            gen_prompt, shot_refs, dropped, shot_mode_eff, shot_tag_driven = _resolve_beat_prompt_refs(
-            if shot_tag_driven:
+                base_prompt,
-                shot_mode_eff = shot_mode
+                ref_slots,
-            else:
+                beat_text,
-                shot_mode_eff = ("every shot" if shot_mode == "auto ref2v"
+                connected_ref_count,
-                                 else "first shot" if shot_mode == "where tagged" else shot_mode)
+                ref_mode,
+                any_tags_anywhere,
+                i,
+                handoff,
+            )
             carry_keyframe = False       # tagged shot keeps its handoff as a keyframe
-            if shot_tag_driven:
+            for n in dropped:
-                # The prompt itself says where each reference belongs: the shot whose
+                if n not in ref_missing:
-                # text names <Picture N> gets image N, renumbered to match what that
+                    ref_missing.append(n)
-                # shot actually carries. Every untagged shot keeps its handoff.
+            context = _reference_context_for_text(
-                gen_prompt, shot_refs, dropped = resolve_tag_driven_prompt_refs(gen_prompt, ref_slots)
+                gen_prompt,
-                for n in dropped:
+                ref_slots,
-                    if n not in ref_missing:
+                resolved_refs=shot_refs,
-                        ref_missing.append(n)
+            )
-            else:
+            if context:
-                shot_refs = shot_references(ref_slots, shot_mode_eff, i, handoff)
+                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

tests/test_dumas_h3_longvideos.py

@@ -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"},