nodes/olivv-cs-ComfyUI-FunPack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

sampler: JoyAI-Echo memory mode (phase 1 — video bank)

Browse Source 
Diagnostic conclusion from the JoyAI-Echo engine: conditioning is fed to the sampler
unchanged (gemma+connector, which Studio already produces), so the ONLY missing sampler
feature is generate_BidirectionalMemory's cross-shot paired AV memory bank. This lands the
VIDEO half: a managed set of clean prior-shot frames injected into each scene so identity
carries across the whole chain — JoyAI's story-level consistency.

Generalizes the single-anchor mid_scene_guide to N frames, reusing the same LTX
guide-attention primitive (_append_guide_latent). _JoyAIMemoryBank holds the policy
(num_fix opening shots pinned as a global anchor + rolling most-recent window capped at
max_size, matching configs/inference.yaml: max_size 7, num_fix_frames 3, center select).
Harvest one clean frame per finished scene; inject the bank before each continuation scene;
supersedes mid_scene_guide when on. Honors the 0.25 audio-safe strength floor.

Off by default, so the audio protect/pass-through invariant is unchanged. Paired audio
memory + v2a_grad_scale coupling is the next phase (needs the ltx_core v2a term).

Inputs: joyai_memory, joyai_memory_size, joyai_fix_frames, joyai_frame_select,
joyai_memory_strength. Threaded through batch-training; logged in run_mechanisms + the
refinement sampler-context. New test_joyai_memory.py covers the bank policy (6 tests).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
2026-06-24 18:16:30 +03:00
co-authored by Claude Opus 4.8
parent 716e21b562
commit 44d12c3203
2 changed files with 185 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
samplers.py
+114 -2
View File
@@ -747,6 +747,37 @@ def _video_only(x_new, x_old, mask):
return x_new
class _JoyAIMemoryBank:
"""JoyAI-Echo cross-shot memory: a rolling set of clean prior-shot latent frames.
Mirrors PairedAudioVideoMemoryBank's policy (JoyAI_Echo configs/inference.yaml): the first
``num_fix`` entries are pinned permanently (the opening shots = global story anchor); the rest
is a most-recent window so the total never exceeds ``max_size``. Stores one latent frame per
finished scene ([B,C,1,H,W]); the audio half rides the existing protect/pass-through path until
the paired-audio phase. Pure bookkeeping — injection is done by the sampler via LTX guide
attention, so this class holds no model state and is unit-testable on its own.
"""
def __init__(self, max_size=7, num_fix=3):
self.max_size = max(1, int(max_size))
self.num_fix = max(0, min(int(num_fix), self.max_size))
self.entries = [] # list of latent-frame tensors, oldest first
def add(self, frame):
if frame is None:
return
self.entries.append(frame)
if len(self.entries) <= self.max_size:
return
fixed = self.entries[:self.num_fix]
recent = self.entries[self.num_fix:]
room = self.max_size - len(fixed)
self.entries = fixed + (recent[-room:] if room > 0 else [])
def frames(self):
return list(self.entries)
def _find_schedule_anchor_index(sigmas, total_steps, schedule_progress):
if sigmas is None or total_steps <= 1:
return 0
@@ -1891,6 +1922,26 @@ class FunPackLTXAVSceneChainSampler:
"default": 0.25, "min": 0.25, "max": 0.5, "step": 0.05,
"tooltip": "Guide attention strength for mid-scene anchor. 0.25 is the minimum — below that audio degrades and character appearance drifts. Above 0.35 causes spatial conflicts when scene composition shifts.",
}),
"joyai_memory": ("BOOLEAN", {
"default": False,
"tooltip": "JoyAI-Echo cross-shot memory bank. Generalizes mid_scene_guide from one anchor to a managed set of clean prior-shot frames injected into each scene via LTX guide attention, so character/scene identity carries across the whole chain (JoyAI-Echo's story-level consistency). The first joyai_fix_frames scenes are pinned permanently as a global anchor; the rest is a rolling most-recent window capped at joyai_memory_size. Supersedes mid_scene_guide when on. Video-only in this phase (audio memory + v2a is a separate opt-in).",
}),
"joyai_memory_size": ("INT", {
"default": 7, "min": 1, "max": 32,
"tooltip": "Max total memory entries injected per scene (JoyAI default 7). Higher = stronger long-range consistency but more guide tokens and slower scenes.",
}),
"joyai_fix_frames": ("INT", {
"default": 3, "min": 0, "max": 16,
"tooltip": "Number of opening scenes pinned permanently in the bank as a global anchor (JoyAI default 3). They are never pruned; entries beyond them are a rolling most-recent window.",
}),
"joyai_frame_select": (["center", "first", "random"], {
"default": "center",
"tooltip": "Which frame of each finished scene to store in the bank (JoyAI default 'center').",
}),
"joyai_memory_strength": ("FLOAT", {
"default": 0.3, "min": 0.25, "max": 0.5, "step": 0.05,
"tooltip": "Guide-attention strength for each memory frame. Same 0.25 floor as mid_scene_guide (below it audio degrades and identity drifts).",
}),
"embed_guidance": ("BOOLEAN", {
"default": False,
"tooltip": "Apply the Refiner's learned quality direction at each denoising step, not just once before sampling. Requires refinement_key_input and enough liked generations to have a direction. Adds ~20-30% inference overhead.",
@@ -3212,6 +3263,45 @@ class FunPackLTXAVSceneChainSampler:
return result, positive, negative, 1
def _harvest_joyai_frame(self, sampled, select):
"""Pick one clean latent frame from a finished scene for the memory bank.
'center' (default) / 'first' / 'random', mirroring JoyAI's frame_selection_mode."""
tensors = self._latent_tensors(sampled)
if not tensors:
return None
F = self._tensor_frames(tensors[0])
if F <= 0:
return None
if select == "first":
idx = 0
elif select == "random":
idx = int(torch.randint(0, F, (1,)).item())
else:
idx = F // 2
return self._time_slice(tensors[0], idx, idx + 1).detach()
def _append_joyai_memory_guides(self, chunk, frames, positive, negative, vae, strength):
"""JoyAI-Echo memory injection: attach every banked prior-shot frame to the current chunk
as a clean LTX guide, so attention carries identity/scene forward across the whole chain.
The single-anchor mid_scene_guide generalized to N frames — reuses _append_guide_latent
per entry. Memory frames are placed at distinct early (prefix) positions, matching JoyAI's
sequence-prefix concat. Audio-safe (guide tokens influence via attention, not overwrite).
Returns the total appended-frame count so the caller crops them off the output tail."""
if not frames:
return chunk, positive, negative, 0
chunk_tensors = self._latent_tensors(chunk)
if not chunk_tensors:
return chunk, positive, negative, 0
F_chunk = self._tensor_frames(chunk_tensors[0])
s = max(0.25, float(strength)) # same audio-safe floor as mid_scene_guide
tail = 0
for i, gf in enumerate(frames):
apply_at = min(max(0, F_chunk - 1), i) # prefix context: distinct early positions
chunk, positive, negative, t = self._append_guide_latent(
chunk, gf, apply_at, s, positive, negative, vae,
)
tail += t
return chunk, positive, negative, tail
def _vae_with_decode_noise(self, vae, timestep, scale, seed):
"""Return a shallow copy of the VAE stamped with LTX decode-time noise settings so its
@@ -3236,6 +3326,8 @@ class FunPackLTXAVSceneChainSampler:
num_frames_per_scene, frame_overlap, cfg, max_scenes, use_same_seed=False,
carry_i2v_guides=False,
mid_scene_guide=False, mid_scene_guide_strength=0.4,
joyai_memory=False, joyai_memory_size=7, joyai_fix_frames=3,
joyai_frame_select="center", joyai_memory_strength=0.3,
embed_guidance=False, embed_guidance_strength=0.02,
embed_guidance_source="relative",
score_slider=False, score_slider_strength=1.0,
@@ -3267,6 +3359,9 @@ class FunPackLTXAVSceneChainSampler:
embed_guidance, embed_guidance_strength, transition_duration,
decode_tile_size, refinement_key_input, embed_guidance_source,
score_slider, score_slider_strength,
joyai_memory=joyai_memory, joyai_memory_size=joyai_memory_size,
joyai_fix_frames=joyai_fix_frames, joyai_frame_select=joyai_frame_select,
joyai_memory_strength=joyai_memory_strength,
)
max_scene_count = max(1, int(max_scenes))
@@ -3329,6 +3424,7 @@ class FunPackLTXAVSceneChainSampler:
scene_outputs: list = []
scene_media_by_ref = self._parse_scene_anchors(funpack_scene_media_refs)
scene_runs: list = []
joyai_bank = _JoyAIMemoryBank(joyai_memory_size, joyai_fix_frames) if joyai_memory else None
for scene_index, scene_cond in enumerate(scene_conditionings):
scene_positive = [scene_cond]
@@ -3404,7 +3500,13 @@ class FunPackLTXAVSceneChainSampler:
chunk, latent_template, scene_positive, scene_negative,
)
carried_guide_frames = max(carried_guide_frames, carried)
if mid_scene_guide and not custom_guides:
if joyai_memory and joyai_bank is not None and not custom_guides:
mem_frames = joyai_bank.frames()
run_mechanisms.append(f"joyai_memory({len(mem_frames)}/{joyai_memory_size},fix={joyai_fix_frames})")
chunk, scene_positive, scene_negative, guide_tail = self._append_joyai_memory_guides(
chunk, mem_frames, scene_positive, scene_negative, vae, joyai_memory_strength,
)
elif mid_scene_guide and not custom_guides:
run_mechanisms.append("mid_scene_guide")
chunk, scene_positive, scene_negative, guide_tail = self._append_mid_scene_guide(
chunk, output, scene_positive, scene_negative, vae, mid_scene_guide_strength,
@@ -3477,6 +3579,10 @@ class FunPackLTXAVSceneChainSampler:
sampled = self._crop_video_head(sampled, carried + soft_carried)
if guide_tail > 0:
sampled = self._crop_video_tail(sampled, guide_tail)
if joyai_bank is not None:
# Harvest from the clean, fully-cropped scene so injected guide tails never re-enter
# the bank. Scene 0 seeds the pinned anchor (num_fix); later scenes roll in.
joyai_bank.add(self._harvest_joyai_frame(sampled, joyai_frame_select))
scene_outputs.append(self._clone_latent(sampled))
blend_overlap = 0 if anchor_meta else video_overlap
output = sampled if output is None else self._blend_latents(output, sampled, blend_overlap)
@@ -3578,6 +3684,7 @@ class FunPackLTXAVSceneChainSampler:
"carry_i2v_guides": bool(carry_i2v_guides),
"frame_overlap": int(frame_overlap),
"transitions_enabled": scene_count > 1,
"joyai_memory": bool(joyai_memory),
})
except Exception as e:
print(f"[FunPackLTXAVSceneChainSampler] Failed to write sampler context: {e}")
@@ -3668,7 +3775,9 @@ class FunPackLTXAVSceneChainSampler:
use_same_seed, carry_i2v_guides, mid_scene_guide, mid_scene_guide_strength,
embed_guidance, embed_guidance_strength, transition_duration,
decode_tile_size, refinement_key_input, embed_guidance_source="relative",
score_slider=False, score_slider_strength=1.0):
score_slider=False, score_slider_strength=1.0,
joyai_memory=False, joyai_memory_size=7, joyai_fix_frames=3,
joyai_frame_select="center", joyai_memory_strength=0.3):
"""Sample one chain per Studio-packed variant entry (seed + index), persisting each result
(latent + preview + per-entry cond + manifest) under ComfyUI temp for rating in Studio.
Reuses sample() per entry with only the seed changed, so each entry is a clean generation."""
@@ -3708,6 +3817,9 @@ class FunPackLTXAVSceneChainSampler:
mid_scene_guide_strength=mid_scene_guide_strength, embed_guidance=embed_guidance,
embed_guidance_strength=embed_guidance_strength, embed_guidance_source=embed_guidance_source,
score_slider=score_slider, score_slider_strength=score_slider_strength,
joyai_memory=joyai_memory, joyai_memory_size=joyai_memory_size,
joyai_fix_frames=joyai_fix_frames, joyai_frame_select=joyai_frame_select,
joyai_memory_strength=joyai_memory_strength,
transition_duration=transition_duration,
decode_tile_size=decode_tile_size, refinement_key_input=key,
unique_id=None, prompt=None,
tests/test_joyai_memory.py
+71
View File
@@ -0,0 +1,71 @@
"""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