Cache H3 prompt analysis helpers
This commit is contained in:
+80
-50
@@ -38,8 +38,9 @@ Verified against ComfyUI core (comfy_extras/nodes_minimax_h3.py, model_base.py,
|
||||
ldm/minimax/model.py, text_encoders/minimax.py, sd.py).
|
||||
"""
|
||||
|
||||
import gc
|
||||
import json
|
||||
import gc
|
||||
from functools import lru_cache
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
@@ -1617,19 +1618,22 @@ def compose_persistent(body, active, anchor_id, removed=None, departed=None,
|
||||
return (count_prefix + out).strip()
|
||||
|
||||
|
||||
def extract_wardrobe(body):
|
||||
def extract_wardrobe(body):
|
||||
"""Pull a 'wardrobe: ...' directive line out of a beat body. Returns
|
||||
(clean_body, wardrobe_or_None). The directive is a whole line starting with
|
||||
'wardrobe:' (case-insensitive), placed INSIDE a beat (not as its own blank-
|
||||
line-separated paragraph, which would become its own shot). It's removed
|
||||
from the body so the literal 'wardrobe:' text isn't stamped as an action."""
|
||||
kept, wardrobe = [], None
|
||||
for ln in body.split("\n"):
|
||||
if re.match(r"\s*wardrobe\s*:", ln, re.I):
|
||||
wardrobe = ln.split(":", 1)[1].strip()
|
||||
else:
|
||||
kept.append(ln)
|
||||
return "\n".join(kept).strip(), wardrobe
|
||||
for ln in body.split("\n"):
|
||||
if re.match(r"\s*wardrobe\s*:", ln, re.I):
|
||||
wardrobe = ln.split(":", 1)[1].strip()
|
||||
else:
|
||||
kept.append(ln)
|
||||
return "\n".join(kept).strip(), wardrobe
|
||||
|
||||
|
||||
extract_wardrobe = lru_cache(maxsize=2048)(extract_wardrobe)
|
||||
|
||||
|
||||
# --- anchor hazards ---------------------------------------------------------
|
||||
@@ -2038,7 +2042,7 @@ _SPOKEN_CUE = re.compile(
|
||||
re.I)
|
||||
|
||||
|
||||
def has_speech(body):
|
||||
def has_speech(body):
|
||||
"""True only if a beat contains ACTUAL scripted speech -- double-quoted words
|
||||
or an explicit <d>...</d> tag. Bare speech VERBS ('calls out', 'tells', 'says'
|
||||
with no quoted line) deliberately do NOT count: unscripted speech is exactly
|
||||
@@ -2062,10 +2066,13 @@ def has_speech(body):
|
||||
lead = body[max(0, m.start() - 60):m.start()].lower()
|
||||
written = [x.end() for x in _WRITTEN_CUE.finditer(lead)]
|
||||
spoken = [x.end() for x in _SPOKEN_CUE.finditer(lead)]
|
||||
if written and (not spoken or written[-1] > spoken[-1]):
|
||||
continue # printed in the scene, nobody said it
|
||||
return True
|
||||
return False
|
||||
if written and (not spoken or written[-1] > spoken[-1]):
|
||||
continue # printed in the scene, nobody said it
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
has_speech = lru_cache(maxsize=2048)(has_speech)
|
||||
|
||||
|
||||
def _spoken_quotes(body):
|
||||
@@ -2709,15 +2716,18 @@ def removed_phrase_items(body, anchor_id):
|
||||
return out
|
||||
|
||||
|
||||
def extract_directive(body, key):
|
||||
def extract_directive(body, key):
|
||||
"""Pull a '<key>: ...' line out of a beat body. Returns (clean_body, value|None)."""
|
||||
kept, val = [], None
|
||||
for ln in body.split("\n"):
|
||||
if re.match(r"\s*" + key + r"\s*:", ln, re.I):
|
||||
val = ln.split(":", 1)[1].strip()
|
||||
else:
|
||||
kept.append(ln)
|
||||
return "\n".join(kept).strip(), val
|
||||
for ln in body.split("\n"):
|
||||
if re.match(r"\s*" + key + r"\s*:", ln, re.I):
|
||||
val = ln.split(":", 1)[1].strip()
|
||||
else:
|
||||
kept.append(ln)
|
||||
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
|
||||
@@ -2860,7 +2870,7 @@ _CLAUSE_SPLIT = (r"(?:[.!?;]+|,?\s+(?:and then|then|and|before|after|while|as|un
|
||||
r"|,\s+(?=[a-z]+ing\b))")
|
||||
|
||||
|
||||
def action_clauses(beat):
|
||||
def action_clauses(beat):
|
||||
"""How many distinct staged actions a beat contains.
|
||||
|
||||
"takes off her red jacket and drops it on the workbench" is two; "walks the
|
||||
@@ -2869,49 +2879,66 @@ def action_clauses(beat):
|
||||
body, _ = extract_wardrobe((beat or "").strip())
|
||||
body = re.sub(r'["“][^"”]*["”]', " ", body)
|
||||
body = " ".join(ln for ln in body.splitlines() if not is_directive_line(ln))
|
||||
parts = [p.strip() for p in re.split(_CLAUSE_SPLIT, body) if p and p.strip()]
|
||||
# A fragment of one word is a leftover ("it", "her"), not an action of its own.
|
||||
return sum(1 for p in parts if len(p.split()) >= 2)
|
||||
parts = [p.strip() for p in re.split(_CLAUSE_SPLIT, body) if p and p.strip()]
|
||||
# A fragment of one word is a leftover ("it", "her"), not an action of its own.
|
||||
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):
|
||||
def estimate_beat_seconds(beat):
|
||||
"""Screen time this beat needs, from its own content. 0.0 when it has none.
|
||||
|
||||
Action and dialogue OVERLAP rather than add -- people talk while they move --
|
||||
so the estimate is the larger of the two, not their sum."""
|
||||
n = action_clauses(beat)
|
||||
action = (BEAT_BASE_SEC + SECONDS_PER_ACTION * n) if n else 0.0
|
||||
return max(action, dialogue_seconds(beat))
|
||||
n = action_clauses(beat)
|
||||
action = (BEAT_BASE_SEC + SECONDS_PER_ACTION * n) if n else 0.0
|
||||
return max(action, dialogue_seconds(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 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_spans(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()]
|
||||
def dialogue_words(beat):
|
||||
"""Words inside double quotes in a beat -- the only speech H3 actually renders."""
|
||||
return sum(_dialogue_spans_cached(beat))
|
||||
|
||||
|
||||
dialogue_words = lru_cache(maxsize=2048)(dialogue_words)
|
||||
|
||||
|
||||
def dialogue_words(beat):
|
||||
"""Words inside double quotes in a beat -- the only speech H3 actually renders."""
|
||||
return sum(dialogue_spans(beat))
|
||||
|
||||
|
||||
def dialogue_seconds(beat, pad=True):
|
||||
def dialogue_seconds(beat, pad=True):
|
||||
"""Screen time this beat's dialogue needs, 0.0 when the beat has none.
|
||||
|
||||
Counts every turn, so a two-character exchange is sized from the WHOLE
|
||||
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)
|
||||
if not spans:
|
||||
return 0.0
|
||||
return (sum(spans) / WORDS_PER_SEC
|
||||
+ TURN_GAP_SEC * (len(spans) - 1)
|
||||
+ (SPEECH_PAD_SEC if pad else 0.0))
|
||||
spans = _dialogue_spans_cached(beat)
|
||||
if not spans:
|
||||
return 0.0
|
||||
return (sum(spans) / WORDS_PER_SEC
|
||||
+ TURN_GAP_SEC * (len(spans) - 1)
|
||||
+ (SPEECH_PAD_SEC if pad else 0.0))
|
||||
|
||||
|
||||
dialogue_seconds = lru_cache(maxsize=4096)(dialogue_seconds)
|
||||
|
||||
|
||||
def beat_seconds_directive(beat):
|
||||
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."""
|
||||
for key in ("seconds", "duration"):
|
||||
@@ -2923,9 +2950,12 @@ def beat_seconds_directive(beat):
|
||||
v = float(m.group(1))
|
||||
except ValueError:
|
||||
continue
|
||||
if v > 0:
|
||||
return v
|
||||
return None
|
||||
if v > 0:
|
||||
return v
|
||||
return None
|
||||
|
||||
|
||||
beat_seconds_directive = lru_cache(maxsize=2048)(beat_seconds_directive)
|
||||
|
||||
|
||||
def plan_beat_frames(beats, fps, budget, per_beat=True):
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user