digital-garbageandClaude Opus 4.8
593267daa8
Extend the JoyAI cross-shot memory bank from video-only to paired audio-video, the part that distinguishes JoyAI-Echo from a generic video-consistency trick. - _JoyAIMemoryBank now stores (video, audio) pairs under the same pin/roll policy; frames()/audio() stay 1:1. Audio half is None until joyai_audio_memory is on, so the phase-1 video path is byte-identical. - _append_joyai_audio_memory pins each banked prior-shot audio latent into the chunk's audio stream as a fully-protected (mask=0) tail, cropped back off after sampling (_crop_audio_tail). The audio analogue of the video guide bank; rides native audio self-attention since LTX has no audio guide-attention API. - _harvest_joyai_audio harvests the paired audio frame, position-matched to video. - v2a_grad_scale scales the model's trained video->audio cross-attention via reversible forward hooks on each AV block's video_to_audio_attn submodule (JoyAI uses 2.0). 1.0 installs no hook -> native behavior, zero overhead. - joyai_audio_memory + v2a_grad_scale wired through sample(), batch training, refinement context, run_mechanisms, INPUT_TYPES (append-only), and editor knobs. Deliberately breaks the audio pass-through invariant, so both are opt-in and off/native by default. 3 new bank tests; zero regressions (same 12 pre-existing stale-mock failures before/after). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
99 lines
3.1 KiB
Python
99 lines
3.1 KiB
Python
"""Unit tests for the JoyAI-Echo cross-shot memory bank policy on the Scene Chain sampler
|
|
(_JoyAIMemoryBank). Covers the fixed-anchor pinning, rolling most-recent window, and the
|
|
max_size cap — the bookkeeping that drives which prior-shot frames get injected each scene.
|
|
"""
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
# Minimal comfy stubs so `import samplers` works without a full ComfyUI env.
|
|
for _name in (
|
|
"comfy", "comfy.k_diffusion", "comfy.k_diffusion.sampling",
|
|
"comfy.model_sampling", "comfy.nested_tensor", "comfy.sample",
|
|
"comfy.samplers", "comfy.utils",
|
|
):
|
|
sys.modules.setdefault(_name, types.ModuleType(_name))
|
|
sys.modules["comfy.nested_tensor"].NestedTensor = object
|
|
|
|
import samplers # noqa: E402
|
|
|
|
Bank = samplers._JoyAIMemoryBank
|
|
|
|
|
|
def test_keeps_everything_under_capacity():
|
|
b = Bank(max_size=7, num_fix=3)
|
|
for i in range(5):
|
|
b.add(i)
|
|
assert b.frames() == [0, 1, 2, 3, 4]
|
|
|
|
|
|
def test_pins_fixed_anchors_and_rolls_recent():
|
|
b = Bank(max_size=5, num_fix=2)
|
|
for i in range(8): # 0..7
|
|
b.add(i)
|
|
# First 2 (the opening shots) pinned forever; the rest is the most-recent window.
|
|
frames = b.frames()
|
|
assert len(frames) == 5
|
|
assert frames[:2] == [0, 1] # anchors never pruned
|
|
assert frames[2:] == [5, 6, 7] # rolling most-recent fills the remaining 3 slots
|
|
|
|
|
|
def test_no_fixed_anchors_is_pure_rolling():
|
|
b = Bank(max_size=3, num_fix=0)
|
|
for i in range(6):
|
|
b.add(i)
|
|
assert b.frames() == [3, 4, 5]
|
|
|
|
|
|
def test_num_fix_clamped_to_max_size():
|
|
b = Bank(max_size=3, num_fix=10) # nonsensical: more anchors than capacity
|
|
for i in range(6):
|
|
b.add(i)
|
|
# num_fix clamps to max_size, so the first 3 pin and nothing rolls.
|
|
assert b.frames() == [0, 1, 2]
|
|
|
|
|
|
def test_add_none_is_ignored():
|
|
b = Bank(max_size=4, num_fix=1)
|
|
b.add(0)
|
|
b.add(None)
|
|
b.add(1)
|
|
assert b.frames() == [0, 1]
|
|
|
|
|
|
def test_frames_returns_a_copy():
|
|
b = Bank(max_size=4, num_fix=1)
|
|
b.add(0)
|
|
out = b.frames()
|
|
out.append(99)
|
|
assert b.frames() == [0] # internal state not mutated by caller
|
|
|
|
|
|
def test_paired_audio_aligns_with_video():
|
|
# Each entry stores a (video, audio) pair; audio() is 1:1 with frames() and shares the policy.
|
|
b = Bank(max_size=3, num_fix=1)
|
|
for i in range(5): # 0..4
|
|
b.add(i, f"a{i}")
|
|
assert b.frames() == [0, 3, 4] # anchor 0 pinned + rolling 3,4
|
|
assert b.audio() == ["a0", "a3", "a4"] # audio pruned in lockstep with video
|
|
|
|
|
|
def test_audio_defaults_to_none_when_video_only():
|
|
# Phase-1 path: add(video) with no audio -> audio() is all None, frames() unaffected.
|
|
b = Bank(max_size=4, num_fix=2)
|
|
b.add(0)
|
|
b.add(1)
|
|
assert b.frames() == [0, 1]
|
|
assert b.audio() == [None, None]
|
|
|
|
|
|
def test_add_none_video_ignores_paired_audio():
|
|
# A None video frame is dropped even if an audio frame was offered (can't pin orphan audio).
|
|
b = Bank(max_size=4, num_fix=1)
|
|
b.add(0, "a0")
|
|
b.add(None, "a1")
|
|
assert b.frames() == [0]
|
|
assert b.audio() == ["a0"]
|