Expand H3 ref2v slots and tag routing

README.md

@@ -34,10 +34,11 @@
   - Reads back the nine optional images for a selected MiniMax H3 plan scene, for example by connecting the current `clip_index`.
 
 - `Dumas H3 Long Videos (FL2VA + REF2VA)`
-  - Inputs: H3 model stack, prompt socket, optional `first_frame`, optional `ref_image_1`..`ref_image_4`, plus the upstream long-video control surface for pacing, continuity, audio, overlays, and guards
+  - Inputs: H3 model stack, prompt socket, optional `first_frame`, optional `ref_image_1`..`ref_image_9`, plus the upstream long-video control surface for pacing, continuity, audio, overlays, and guards
   - Outputs: `images`, `audio`, `info`, `script`, `frames_per_shot`, `total_frames`, `shots`, `video_seconds`, `fps`, `fps_int`, `latent`, `soundscape`
   - First-pass Dumas port of the `MiniMax-H3-Longvideos` sampler, brought in as a local starting point for long-form H3 chaining work.
   - Keeps the upstream split-beats / handoff / ref-routing behavior close to source so future Dumas-specific improvements can be compared against a known baseline.
+  - Prompt `<Picture N>` tags now map to the actual ref socket numbers you wire, even with gaps such as only `ref_image_2` and `ref_image_7` connected.
 
 - `Dumas H3 Shot Length`
   - Inputs: `shot_seconds`, `fps`, optional `cap_to_h3_max`

dumas_h3_longvideos.py

@@ -255,16 +255,17 @@ def split_paragraphs(text, delimiter):
 # INPUT_TYPES so a workflow saved before they existed still maps its stored values
 # onto the right widgets (ComfyUI matches them by position, not by name).
 # APPEND to this tuple when adding a widget; never insert into the middle.
-ADDED_WIDGETS = (
+ADDED_WIDGETS = (
-    "beat_split", "per_beat_length",
+    "beat_split", "per_beat_length",
-    "watermark_text", "watermark_position", "watermark_size", "watermark_opacity",
+    "watermark_text", "watermark_position", "watermark_size", "watermark_opacity",
-    "watermark_margin", "intro_text", "intro_position", "intro_seconds",
+    "watermark_margin", "intro_text", "intro_position", "intro_seconds",
-    "intro_fade", "intro_size", "overlay_font", "overlay_stroke",
+    "intro_fade", "intro_size", "overlay_font", "overlay_stroke",
-    "ref_mode", "ref_image_size", "ref_noise_aug", "auto_props", "prevent_nudity",
+    "ref_mode", "ref_image_size", "ref_noise_aug", "auto_props", "prevent_nudity",
-    "exposed_terms", "anatomy_guard", "lock_restraints", "solidity_guard",
+    "exposed_terms", "anatomy_guard", "lock_restraints", "solidity_guard",
-    "motion_guard", "contact_guard",
+    "motion_guard", "contact_guard",
-    "auto_soundscape", "allow_nonspeech_vocals",
+    "auto_soundscape", "allow_nonspeech_vocals",
-)
+    "ref_image_5", "ref_image_6", "ref_image_7", "ref_image_8", "ref_image_9",
+)
 
 NL = "\n"
 # Lines that CONFIGURE a beat rather than being one. They attach to the beat that
@@ -3992,26 +3993,31 @@ def keyframe_rides_with_refs(ref_noise_aug):
 _PICTURE_TAG = re.compile(r"<\s*picture[\s_\-]*(\d+)\s*>", re.I)
 
 
-def picture_tags(text):
+def picture_tags(text):
-    """The reference slots a shot's text asks for, in ascending order."""
+    """The reference slots a shot's text asks for, in ascending order."""
-    return sorted({int(m.group(1)) for m in _PICTURE_TAG.finditer(text or "")})
+    return sorted({int(m.group(1)) for m in _PICTURE_TAG.finditer(text or "")})
-
+
-
+
-def resolve_tagged_refs(text, ref_list):
+def _connected_refs(ref_slots):
-    """(rewritten text, images, dropped) for the <Picture N> tags in ONE shot.
+    """Connected refs only, preserving slot order and skipping empty sockets."""
-
+    return [ref for ref in (ref_slots or []) if ref is not None]
-    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
+
+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
+    a shot that uses only <Picture 2> would receive that image labelled
     <Picture 1> and the text would point at nothing. The tags are therefore
     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>.
-
+
-    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 [])]
+    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]
+    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):
         n = int(m.group(1))
@@ -4025,7 +4031,7 @@ def resolve_tagged_refs(text, ref_list):
     return out.strip(), [ref_list[n - 1] for n in live], dropped
 
 
-def shot_references(ref_list, ref_mode, shot_index, handoff):
+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.
 
@@ -4043,19 +4049,20 @@ def shot_references(ref_list, ref_mode, shot_index, handoff):
                        continuous take.
       'every shot + handoff ref'
                     -- every shot is ref-conditioned AND the previous shot's last
-                       frame is appended as one more reference. Continuity comes
+                       frame is appended as one more reference. Continuity comes
-                       back as a soft signal (the model is shown where the last
+                       back as a soft signal (the model is shown where the last
-                       shot ended rather than told to start exactly there), and it
+                       shot ended rather than told to start exactly there), and it
-                       stays a single ref2va task, so nothing conflicts."""
+                       stays a single ref2va task, so nothing conflicts."""
-    if not ref_list:
+    refs = _connected_refs(ref_list)
-        return []
+    if not refs:
-    if ref_mode == "first shot":
+        return []
-        return list(ref_list) if shot_index == 0 else []
+    if ref_mode == "first shot":
-    if ref_mode == "every shot":
+        return list(refs) if shot_index == 0 else []
-        return list(ref_list)
+    if ref_mode == "every shot":
-    if ref_mode == "every shot + handoff ref":
+        return list(refs)
-        return list(ref_list) + ([handoff] if handoff is not None else [])
+    if ref_mode == "every shot + handoff ref":
-    return list(ref_list) if shot_index == 0 else []       # unknown value -> safest
+        return list(refs) + ([handoff] if handoff is not None else [])
+    return list(refs) if shot_index == 0 else []       # unknown value -> safest
 
 
 # --- text-encoder / DiT compatibility -------------------------------------
@@ -5239,18 +5246,27 @@ class H3LongVideos:
             },
             "optional": {
                 "first_frame": ("IMAGE",),
-                # ref2va inputs. Order matters and is the ONLY thing that decides the
+                # ref2va inputs. SOCKET NUMBER matters: prompt tags refer to these
-                # roster: the tokenizer labels these <Picture 1>..<Picture 4> in the
+                # exact slots, even when some intermediate sockets are left empty.
-                # order they appear here, then appends the prompt. Refer to them by
+                # A shot using only ref_image_7 is still tagged as <Picture 7> in the
-                # those tags in the prompt if you want a reference bound to a named
+                # prompt and renumbered only for the per-shot tokenizer payload.
-                # character ("Kristy, <Picture 1>, walks in").
+                # When a mode uses all connected refs, they keep socket order.
-                "ref_image_1": ("IMAGE", {"tooltip": "Reference image <Picture 1> -- identity/appearance "
+                # The tokenizer labels the carried refs <Picture 1>..<Picture N> in
-                    "carried into the shots. Which shots receive it is set by ref_mode (or <Picture N> "
+                # the order they are handed that shot.
-                    "tags in the beats); a referenced shot ALSO carries the previous frame as its "
+                # Refer to socket tags in the prompt if you want a reference bound to
-                    "keyframe, so taking a reference never costs continuity."}),
+                # a named character ("Kristy, <Picture 7>, walks in").
-                "ref_image_2": ("IMAGE", {"tooltip": "Reference image <Picture 2>."}),
+                "ref_image_1": ("IMAGE", {"tooltip": "Reference image <Picture 1> -- identity/appearance "
-                "ref_image_3": ("IMAGE", {"tooltip": "Reference image <Picture 3>."}),
+                    "carried into the shots. Which shots receive it is set by ref_mode (or <Picture N> "
-                "ref_image_4": ("IMAGE", {"tooltip": "Reference image <Picture 4>."}),
+                    "tags in the beats); a referenced shot ALSO carries the previous frame as its "
+                    "keyframe, so taking a reference never costs continuity."}),
+                "ref_image_2": ("IMAGE", {"tooltip": "Reference image <Picture 2>."}),
+                "ref_image_3": ("IMAGE", {"tooltip": "Reference image <Picture 3>."}),
+                "ref_image_4": ("IMAGE", {"tooltip": "Reference image <Picture 4>."}),
+                "ref_image_5": ("IMAGE", {"tooltip": "Reference image <Picture 5>."}),
+                "ref_image_6": ("IMAGE", {"tooltip": "Reference image <Picture 6>."}),
+                "ref_image_7": ("IMAGE", {"tooltip": "Reference image <Picture 7>."}),
+                "ref_image_8": ("IMAGE", {"tooltip": "Reference image <Picture 8>."}),
+                "ref_image_9": ("IMAGE", {"tooltip": "Reference image <Picture 9>."}),
                 "plan_only": ("BOOLEAN", {"default": False,
                     "tooltip": "Preview the shot split WITHOUT rendering. Uses THIS node's own settings (no "
                                "second node, no duplicate entry): returns the plan in 'info' and the "
@@ -5743,11 +5759,13 @@ class H3LongVideos:
             mute_nonspeech_audio=True, mute_fade_ms=40,
             watermark_text="", watermark_position="bottom-right", watermark_size=4.0,
             watermark_opacity=0.75, watermark_margin=3.0,
-            intro_text="", intro_position="center", intro_seconds=3.0, intro_fade=0.6,
+            intro_text="", intro_position="center", intro_seconds=3.0, intro_fade=0.6,
-            intro_size=9.0, overlay_font="arial.ttf", overlay_stroke=0,
+            intro_size=9.0, overlay_font="arial.ttf", overlay_stroke=0,
-            ref_image_1=None, ref_image_2=None, ref_image_3=None, ref_image_4=None,
+            ref_image_1=None, ref_image_2=None, ref_image_3=None, ref_image_4=None,
-            ref_mode="where tagged", ref_image_size="match", ref_noise_aug=0.999,
+            ref_image_5=None, ref_image_6=None, ref_image_7=None, ref_image_8=None,
-            graph=None, node_id=None):
+            ref_image_9=None,
+            ref_mode="where tagged", ref_image_size="match", ref_noise_aug=0.999,
+            graph=None, node_id=None):
 
         # FIRST: detect a checkpoint swap since the previous execution and hard-flush.
         # A stale resident model from a different checkpoint would otherwise poison
@@ -6024,21 +6042,24 @@ class H3LongVideos:
                                      else "prompt/soundscape silencing only"))
             # Same reference accounting the render reports: which shots lose the
             # handoff is a composition decision, so it belongs in the preview.
-            n_refs = len([r for r in (ref_image_1, ref_image_2, ref_image_3, ref_image_4)
+            ref_slots = [
-                          if r is not None])
+                ref_image_1, ref_image_2, ref_image_3, ref_image_4, ref_image_5,
+                ref_image_6, ref_image_7, ref_image_8, ref_image_9,
+            ]
+            n_refs = len(_connected_refs(ref_slots))
             plan_ref = ""
             if n_refs:
                 # Mirror the render's placement exactly: 'where tagged' reads the
-                # prompts and falls back to first shot when nothing is tagged --
+                # prompts and falls back to first shot when nothing is tagged --
-                # reporting by ref_mode alone described shots the render never gave
+                # reporting by ref_mode alone described shots the render never gave
-                # references to.
+                # references to.
-                if ref_mode == "where tagged" and any(picture_tags(g) for g in gens):
+                if ref_mode == "where tagged" and any(picture_tags(g) for g in gens):
-                    on = [n + 1 for n, g in enumerate(gens) if picture_tags(g)]
+                    on = [n + 1 for n, g in enumerate(gens) if resolve_tagged_refs(g, ref_slots)[1]]
-                    how = "placed by <Picture N> tags"
+                    how = "placed by <Picture N> tags"
-                else:
+                else:
                     mode_eff = "first shot" if ref_mode == "where tagged" else ref_mode
-                    on = [n + 1 for n in range(shots)
+                    on = [n + 1 for n in range(shots)
-                          if shot_references([1] * n_refs, mode_eff, n, 1 if n else None)]
+                          if shot_references(ref_slots, mode_eff, n, 1 if n else None)]
                     how = (f"ref_mode '{mode_eff}'"
                            + (" -- no tags found anywhere" if ref_mode == "where tagged" else ""))
                 plan_ref = (f" ref2va: {n_refs} reference image(s) at '{ref_image_size}' on shot(s) "
@@ -6076,18 +6097,22 @@ class H3LongVideos:
         latent_chunks = []                 # per-shot sampled latents, pre-decode
         mouth_settled = []                 # shots seeded from a settled (closed) mouth
         handoff, sr = first_frame, None
-        ref_list = [r for r in (ref_image_1, ref_image_2, ref_image_3, ref_image_4) if r is not None]
+        ref_list = [
-        ref_shots = []                     # which shots ended up ref-conditioned
+            ref_image_1, ref_image_2, ref_image_3, ref_image_4, ref_image_5,
-        ref_missing = []                   # <Picture N> tags naming an unconnected slot
+            ref_image_6, ref_image_7, ref_image_8, ref_image_9,
-        ref_carried = []                   # tagged shots that kept continuity as an extra ref
+        ]
-        ref_keyframed = []                 # tagged shots that kept it as a real keyframe
+        connected_ref_count = len(_connected_refs(ref_list))
-        # 'where tagged' reads the prompt instead of counting shots. If references are
+        ref_shots = []                     # which shots ended up ref-conditioned
-        # connected but nothing is tagged anywhere, fall back to first-shot placement
+        ref_missing = []                   # <Picture N> tags naming an unconnected slot
-        # rather than silently conditioning nothing at all.
+        ref_carried = []                   # tagged shots that kept continuity as an extra ref
-        tag_driven = bool(ref_list) and ref_mode == "where tagged" and any(
+        ref_keyframed = []                 # tagged shots that kept it as a real keyframe
-            picture_tags(g) for g in gens)
+        # 'where tagged' reads the prompt instead of counting shots. If references are
-        if ref_list and ref_mode == "where tagged" and not tag_driven:
+        # connected but nothing is tagged anywhere, fall back to first-shot placement
-            ref_mode = "first shot"
+        # rather than silently conditioning nothing at all.
+        tag_driven = bool(connected_ref_count) and ref_mode == "where tagged" and any(
+            picture_tags(g) for g in gens)
+        if connected_ref_count and ref_mode == "where tagged" and not tag_driven:
+            ref_mode = "first shot"
         if cleanup_between_shots:
             _deep_cleanup()          # start the first (heaviest) shot with max free VRAM
 
@@ -6384,10 +6409,10 @@ class H3LongVideos:
                                 f"dropped from the text")
         else:
             ref_note_missing = ""
-        if ref_list and ref_shots:
+        if connected_ref_count and ref_shots:
             kept = [n for n in range(1, len(gens) + 1) if n not in ref_shots]
             ref_placement = "placed by <Picture N> tags" if tag_driven else f"ref_mode '{ref_mode}'"
-            ref_note = (f" ref2va: {len(ref_list)} reference image(s) at '{ref_image_size}' on shot(s) "
+            ref_note = (f" ref2va: {connected_ref_count} reference image(s) at '{ref_image_size}' on shot(s) "
                         f"{','.join(str(n) for n in ref_shots)} "
                         f"({ref_placement})"
                         + (f", ref_noise_aug {ref_noise_aug:.3f}" if ref_noise_aug is not None
@@ -6401,14 +6426,14 @@ class H3LongVideos:
                            f"{','.join(str(n) for n in ref_carried)} -- weaker than a keyframe, but "
                            f"ref_noise_aug below {KEYFRAME_SAFE_AUG:g} would soften a keyframe too "
                            f"(one aug covers every cond latent)" if ref_carried else "")
-                        + ("" if (ref_keyframed or ref_carried or kept)
+                        + ("" if (ref_keyframed or ref_carried or kept)
-                           else ", so every cut between beats is a CUT, not a continuous take")
+                           else ", so every cut between beats is a CUT, not a continuous take")
-                        + ref_note_missing)
+                        + ref_note_missing)
-        elif ref_list:
+        elif connected_ref_count:
-            ref_note = (f" ref2va: {len(ref_list)} reference image(s) connected but ref_mode "
+            ref_note = (f" ref2va: {connected_ref_count} reference image(s) connected but ref_mode "
-                        f"'{ref_mode}' applied them to no shot")
+                        f"'{ref_mode}' applied them to no shot")
-        else:
+        else:
-            ref_note = ""
+            ref_note = ""
         info = ((anchor_note + " ") if anchor_note else "") + \
                (f"{shape_str} at {w}x{h}; {all_frames.shape[0]} frames (~{actual:.1f}s actual). "
                 f"decode {'tiled' if tiled else 'full'}. {vram_str}.{hoff_str}"

tests/test_dumas_h3_longvideos.py

@@ -132,6 +132,51 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
         fn(spoken)
         self.assertGreaterEqual(fn.cache_info().hits, 2)
 
+    def test_resolve_tagged_refs_preserves_sparse_socket_numbers(self):
+        refs = [None, "img2", None, None, None, None, "img7", None, "img9"]
+
+        text, images, dropped = self.module.resolve_tagged_refs(
+            "Mara <Picture 7> turns toward Jon <Picture 2> while <Picture 9> watches.",
+            refs,
+        )
+
+        self.assertEqual(
+            text,
+            "Mara <Picture 2> turns toward Jon <Picture 1> while <Picture 3> watches.",
+        )
+        self.assertEqual(images, ["img2", "img7", "img9"])
+        self.assertEqual(dropped, [])
+
+    def test_resolve_tagged_refs_drops_unconnected_sparse_slots(self):
+        refs = [None, "img2", None, None, None, None, "img7", None, None]
+
+        text, images, dropped = self.module.resolve_tagged_refs(
+            "Use <Picture 7>, skip <Picture 4>, keep <Picture 2>.",
+            refs,
+        )
+
+        self.assertEqual(text, "Use <Picture 2>, skip, keep <Picture 1>.")
+        self.assertEqual(images, ["img2", "img7"])
+        self.assertEqual(dropped, [4])
+
+    def test_shot_references_uses_all_connected_sparse_slots(self):
+        refs = [None, "img2", None, "img4", None, None, "img7", None, None]
+
+        self.assertEqual(
+            self.module.shot_references(refs, "first shot", 0, None),
+            ["img2", "img4", "img7"],
+        )
+        self.assertEqual(
+            self.module.shot_references(refs, "every shot", 3, None),
+            ["img2", "img4", "img7"],
+        )
+
+    def test_input_types_expose_nine_ref_slots(self):
+        optional = self.module.H3LongVideos.INPUT_TYPES()["optional"]
+
+        for index in range(1, 10):
+            self.assertIn(f"ref_image_{index}", optional)
+
 
 if __name__ == "__main__":
     unittest.main()