Cache H3 prompt analysis helpers

dumas_h3_longvideos.py

@@ -39,6 +39,7 @@ ldm/minimax/model.py, text_encoders/minimax.py, sd.py).
 """
 
 import gc
+from functools import lru_cache
 import json
 import logging
 import math
@@ -1632,6 +1633,9 @@ def extract_wardrobe(body):
     return "\n".join(kept).strip(), wardrobe
 
 
+extract_wardrobe = lru_cache(maxsize=2048)(extract_wardrobe)
+
+
 # --- anchor hazards ---------------------------------------------------------
 # The anchor is stamped into EVERY shot, so anything in it has to be true of every
 # shot. Four kinds of thing are not, and each fails in its own way.
@@ -2068,6 +2072,9 @@ def has_speech(body):
     return False
 
 
+has_speech = lru_cache(maxsize=2048)(has_speech)
+
+
 def _spoken_quotes(body):
     """Each double-quoted span that IS speech, skipping printed text -- the same
     nearest-cue rule has_speech uses, but returning the spans so each one can be
@@ -2720,6 +2727,9 @@ def extract_directive(body, key):
     return "\n".join(kept).strip(), val
 
 
+extract_directive = lru_cache(maxsize=4096)(extract_directive)
+
+
 # "walks out OF THE BARN" is emerging INTO the scene, not leaving it -- and a false
 # exit is the expensive error: the character is stripped from every later shot and
 # only an explicit 'enter:' brings them back. So "out of <somewhere>" is never an
@@ -2874,6 +2884,9 @@ def action_clauses(beat):
     return sum(1 for p in parts if len(p.split()) >= 2)
 
 
+action_clauses = lru_cache(maxsize=2048)(action_clauses)
+
+
 def estimate_beat_seconds(beat):
     """Screen time this beat needs, from its own content. 0.0 when it has none.
 
@@ -2884,16 +2897,27 @@ def estimate_beat_seconds(beat):
     return max(action, dialogue_seconds(beat))
 
 
-def dialogue_spans(beat):
+estimate_beat_seconds = lru_cache(maxsize=2048)(estimate_beat_seconds)
+
+
+@lru_cache(maxsize=2048)
+def _dialogue_spans_cached(beat):
     """Word count of each double-quoted span in a beat, in order. Length of the
     returned list is the number of speaking TURNS -- the multi-character case."""
     body, _ = extract_wardrobe((beat or "").strip())
-    return [len(q.split()) for q in re.findall(r'["\u201c]([^"\u201d]+)["\u201d]', body) if q.split()]
+    return tuple(len(q.split()) for q in re.findall(r'["\u201c]([^"\u201d]+)["\u201d]', body) if q.split())
+
+
+def dialogue_spans(beat):
+    return list(_dialogue_spans_cached(beat))
 
 
 def dialogue_words(beat):
     """Words inside double quotes in a beat -- the only speech H3 actually renders."""
-    return sum(dialogue_spans(beat))
+    return sum(_dialogue_spans_cached(beat))
+
+
+dialogue_words = lru_cache(maxsize=2048)(dialogue_words)
 
 
 def dialogue_seconds(beat, pad=True):
@@ -2903,7 +2927,7 @@ def dialogue_seconds(beat, pad=True):
     exchange plus a gap between turns -- not from the longest single line.
     `pad` controls only the head/tail air; turn gaps are always counted because
     they are time the shot genuinely has to contain."""
-    spans = dialogue_spans(beat)
+    spans = _dialogue_spans_cached(beat)
     if not spans:
         return 0.0
     return (sum(spans) / WORDS_PER_SEC
@@ -2911,6 +2935,9 @@ def dialogue_seconds(beat, pad=True):
             + (SPEECH_PAD_SEC if pad else 0.0))
 
 
+dialogue_seconds = lru_cache(maxsize=4096)(dialogue_seconds)
+
+
 def beat_seconds_directive(beat):
     """Explicit per-beat length: a 'seconds: 8' (or 'duration: 8') line in the beat.
     Returns the float, or None when the beat doesn't set one."""
@@ -2928,6 +2955,9 @@ def beat_seconds_directive(beat):
     return None
 
 
+beat_seconds_directive = lru_cache(maxsize=2048)(beat_seconds_directive)
+
+
 def plan_beat_frames(beats, fps, budget, per_beat=True):
     """Per-beat shot lengths in frames. Returns (lengths, notes).
 

tests/test_dumas_h3_longvideos.py

@@ -0,0 +1,137 @@
+import importlib
+import sys
+import types
+import unittest
+
+
+class DumasH3LongVideosHelperTests(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls._saved_modules = {
+            name: sys.modules.get(name)
+            for name in (
+                "torch",
+                "nodes",
+                "comfy",
+                "comfy.utils",
+                "comfy.samplers",
+                "comfy.nested_tensor",
+                "comfy.model_management",
+                "node_helpers",
+                "dumas_h3_longvideos",
+            )
+        }
+
+        fake_torch = types.SimpleNamespace(
+            cuda=types.SimpleNamespace(OutOfMemoryError=RuntimeError),
+            float32="float32",
+        )
+        fake_nodes = types.SimpleNamespace(common_ksampler=lambda *args, **kwargs: ({},))
+        fake_comfy_samplers = types.SimpleNamespace(
+            KSampler=types.SimpleNamespace(SAMPLERS=("res_multistep",), SCHEDULERS=("simple",))
+        )
+        fake_comfy_utils = types.SimpleNamespace(ProgressBar=lambda total: None)
+        fake_mm = types.SimpleNamespace(
+            current_loaded_models=[],
+            free_memory=lambda *args, **kwargs: None,
+            get_torch_device=lambda: "cpu",
+            soft_empty_cache=lambda *args, **kwargs: None,
+            unload_all_models=lambda *args, **kwargs: None,
+            get_free_memory=lambda *args, **kwargs: 0,
+            get_total_memory=lambda *args, **kwargs: 0,
+        )
+        fake_comfy = types.SimpleNamespace(
+            utils=fake_comfy_utils,
+            samplers=fake_comfy_samplers,
+            nested_tensor=types.SimpleNamespace(),
+            model_management=fake_mm,
+        )
+
+        sys.modules["torch"] = fake_torch
+        sys.modules["nodes"] = fake_nodes
+        sys.modules["comfy"] = fake_comfy
+        sys.modules["comfy.utils"] = fake_comfy_utils
+        sys.modules["comfy.samplers"] = fake_comfy_samplers
+        sys.modules["comfy.nested_tensor"] = fake_comfy.nested_tensor
+        sys.modules["comfy.model_management"] = fake_mm
+        sys.modules["node_helpers"] = types.SimpleNamespace()
+
+        cls.module = importlib.import_module("dumas_h3_longvideos")
+
+    @classmethod
+    def tearDownClass(cls):
+        for name, module in cls._saved_modules.items():
+            if module is None:
+                sys.modules.pop(name, None)
+            else:
+                sys.modules[name] = module
+
+    def test_extract_wardrobe_is_cached(self):
+        fn = self.module.extract_wardrobe
+        fn.cache_clear()
+
+        beat = "walks forward\nwardrobe: red jacket, grey shorts\nlooks back"
+        self.assertEqual(fn(beat), ("walks forward\nlooks back", "red jacket, grey shorts"))
+        self.assertEqual(fn(beat), ("walks forward\nlooks back", "red jacket, grey shorts"))
+        self.assertGreater(fn.cache_info().hits, 0)
+
+    def test_dialogue_helpers_keep_existing_outputs_and_cache(self):
+        spans_cache = self.module._dialogue_spans_cached
+        sec_fn = self.module.dialogue_seconds
+        words_fn = self.module.dialogue_words
+
+        spans_cache.cache_clear()
+        sec_fn.cache_clear()
+        words_fn.cache_clear()
+
+        beat = 'Mara says, "Open it now." Jon replies, "Do it."'
+        self.assertEqual(self.module.dialogue_spans(beat), [3, 2])
+        self.assertEqual(words_fn(beat), 5)
+        self.assertAlmostEqual(sec_fn(beat), 3.5)
+        self.assertAlmostEqual(sec_fn(beat, pad=False), 2.5)
+
+        self.module.dialogue_spans(beat)
+        sec_fn(beat)
+        words_fn(beat)
+
+        self.assertGreater(spans_cache.cache_info().hits, 0)
+        self.assertGreater(sec_fn.cache_info().hits, 0)
+        self.assertGreater(words_fn.cache_info().hits, 0)
+
+    def test_directive_and_estimate_helpers_are_cached(self):
+        directive_fn = self.module.beat_seconds_directive
+        estimate_fn = self.module.estimate_beat_seconds
+        action_fn = self.module.action_clauses
+
+        directive_fn.cache_clear()
+        estimate_fn.cache_clear()
+        action_fn.cache_clear()
+
+        beat = 'seconds: 7.5\nShe opens the hatch and climbs inside.'
+        self.assertEqual(directive_fn(beat), 7.5)
+        self.assertEqual(action_fn(beat), 2)
+        self.assertAlmostEqual(estimate_fn(beat), 7.0)
+
+        directive_fn(beat)
+        action_fn(beat)
+        estimate_fn(beat)
+
+        self.assertGreater(directive_fn.cache_info().hits, 0)
+        self.assertGreater(action_fn.cache_info().hits, 0)
+        self.assertGreater(estimate_fn.cache_info().hits, 0)
+
+    def test_has_speech_cache_respects_written_text_filter(self):
+        fn = self.module.has_speech
+        fn.cache_clear()
+
+        written = 'She reads the sign marked "EXIT" and keeps walking.'
+        spoken = 'She says, "Exit now." and points to the door.'
+        self.assertFalse(fn(written))
+        self.assertTrue(fn(spoken))
+        fn(written)
+        fn(spoken)
+        self.assertGreaterEqual(fn.cache_info().hits, 2)
+
+
+if __name__ == "__main__":
+    unittest.main()