Each mixed scene is its own queue request (max_scenes=1), not continuous chain overlap. Scene 2+ uses its anchor via graph Img2Video and borrows hidden guides from the previous scene. Sampler applies guide stacks on the first chunk for solo mixed runs.
559 lines
18 KiB
Python
559 lines
18 KiB
Python
import sys
|
|
import types
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
|
|
class FakeNestedTensor:
|
|
def __init__(self, tensors):
|
|
self.tensors = list(tensors)
|
|
self.is_nested = True
|
|
|
|
def unbind(self):
|
|
return self.tensors
|
|
|
|
@property
|
|
def shape(self):
|
|
return self.tensors[0].shape
|
|
|
|
@property
|
|
def device(self):
|
|
return self.tensors[0].device
|
|
|
|
@property
|
|
def dtype(self):
|
|
return self.tensors[0].dtype
|
|
|
|
@property
|
|
def layout(self):
|
|
return self.tensors[0].layout
|
|
|
|
def size(self):
|
|
return self.tensors[0].size()
|
|
|
|
|
|
sample_calls = []
|
|
|
|
|
|
def _zeros_like(value):
|
|
if getattr(value, "is_nested", False):
|
|
return FakeNestedTensor([torch.zeros_like(t) for t in value.unbind()])
|
|
return torch.zeros_like(value)
|
|
|
|
|
|
def _sample_like(value, mask, seed):
|
|
if getattr(value, "is_nested", False):
|
|
masks = mask.unbind() if getattr(mask, "is_nested", False) else [None] * len(value.unbind())
|
|
return FakeNestedTensor([
|
|
_sample_like(tensor, masks[index], seed)
|
|
for index, tensor in enumerate(value.unbind())
|
|
])
|
|
if mask is None:
|
|
return value + float(seed)
|
|
return value + mask.to(value.device, value.dtype) * float(seed)
|
|
|
|
|
|
def fake_prepare_noise(samples, seed, noise_inds=None):
|
|
return _zeros_like(samples)
|
|
|
|
|
|
def fake_sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image,
|
|
noise_mask=None, callback=None, disable_pbar=False, seed=None):
|
|
sample_calls.append({
|
|
"seed": seed,
|
|
"positive": positive,
|
|
"negative": negative,
|
|
"cfg": cfg,
|
|
"latent_image": _sample_snapshot(latent_image),
|
|
"noise_mask": _sample_snapshot(noise_mask),
|
|
})
|
|
return _sample_like(latent_image, noise_mask, seed)
|
|
|
|
|
|
def _sample_snapshot(value):
|
|
if value is None:
|
|
return None
|
|
if getattr(value, "is_nested", False):
|
|
return FakeNestedTensor([tensor.detach().clone() for tensor in value.unbind()])
|
|
return value.detach().clone()
|
|
|
|
|
|
comfy_mod = types.ModuleType("comfy")
|
|
comfy_kd_mod = types.ModuleType("comfy.k_diffusion")
|
|
comfy_kd_sampling_mod = types.ModuleType("comfy.k_diffusion.sampling")
|
|
comfy_model_sampling_mod = types.ModuleType("comfy.model_sampling")
|
|
comfy_nested_mod = types.ModuleType("comfy.nested_tensor")
|
|
comfy_sample_mod = types.ModuleType("comfy.sample")
|
|
comfy_samplers_mod = types.ModuleType("comfy.samplers")
|
|
comfy_utils_mod = types.ModuleType("comfy.utils")
|
|
|
|
comfy_nested_mod.NestedTensor = FakeNestedTensor
|
|
comfy_sample_mod.prepare_noise = fake_prepare_noise
|
|
comfy_sample_mod.sample_custom = fake_sample_custom
|
|
|
|
comfy_mod.k_diffusion = comfy_kd_mod
|
|
comfy_kd_mod.sampling = comfy_kd_sampling_mod
|
|
comfy_mod.model_sampling = comfy_model_sampling_mod
|
|
comfy_mod.nested_tensor = comfy_nested_mod
|
|
comfy_mod.sample = comfy_sample_mod
|
|
comfy_mod.samplers = comfy_samplers_mod
|
|
comfy_mod.utils = comfy_utils_mod
|
|
|
|
sys.modules.setdefault("comfy", comfy_mod)
|
|
sys.modules.setdefault("comfy.k_diffusion", comfy_kd_mod)
|
|
sys.modules.setdefault("comfy.k_diffusion.sampling", comfy_kd_sampling_mod)
|
|
sys.modules.setdefault("comfy.model_sampling", comfy_model_sampling_mod)
|
|
sys.modules.setdefault("comfy.nested_tensor", comfy_nested_mod)
|
|
sys.modules.setdefault("comfy.sample", comfy_sample_mod)
|
|
sys.modules.setdefault("comfy.samplers", comfy_samplers_mod)
|
|
sys.modules.setdefault("comfy.utils", comfy_utils_mod)
|
|
|
|
from samplers import FunPackLTXAVSceneChainSampler
|
|
|
|
|
|
class FakeVAE:
|
|
downscale_index_formula = (1, 1, 1)
|
|
|
|
|
|
def scene_cond(index):
|
|
return (
|
|
torch.ones(1, 2, 3) * float(index + 1),
|
|
{"funpack_scene_text": f"scene {index + 1}"},
|
|
)
|
|
|
|
|
|
def test_scene_chain_detects_scene_count_and_increments_seed():
|
|
sample_calls.clear()
|
|
node = FunPackLTXAVSceneChainSampler()
|
|
latent_template = {"samples": torch.zeros(1, 2, 5, 3, 3)}
|
|
positive = [scene_cond(0), scene_cond(1), scene_cond(2)]
|
|
negative = [(torch.zeros(1, 2, 3), {})]
|
|
|
|
latent, status, scene_count, report = node.sample(
|
|
model=object(),
|
|
vae=FakeVAE(),
|
|
positive=positive,
|
|
negative=negative,
|
|
sampler=object(),
|
|
sigmas=torch.tensor([1.0, 0.0]),
|
|
seed=10,
|
|
latent_template=latent_template,
|
|
num_frames_per_scene=5,
|
|
frame_overlap=2,
|
|
cfg=1.5,
|
|
max_scenes=8,
|
|
)
|
|
|
|
assert scene_count == 3
|
|
assert [call["seed"] for call in sample_calls] == [10, 11, 12]
|
|
assert [call["positive"][0][1]["funpack_scene_text"] for call in sample_calls] == ["scene 1", "scene 2", "scene 3"]
|
|
assert latent["samples"].shape[2] == 11
|
|
assert "Scene chain complete" in status
|
|
assert "Scene 3" in report
|
|
|
|
|
|
def test_scene_chain_accepts_manual_combined_conditioning_without_metadata():
|
|
sample_calls.clear()
|
|
node = FunPackLTXAVSceneChainSampler()
|
|
latent_template = {"samples": torch.zeros(1, 2, 5, 3, 3)}
|
|
positive = [
|
|
(torch.ones(1, 2, 3), {}),
|
|
(torch.ones(1, 2, 3) * 2.0, {}),
|
|
(torch.ones(1, 2, 3) * 3.0, {}),
|
|
]
|
|
|
|
_, status, scene_count, report = node.sample(
|
|
model=object(),
|
|
vae=FakeVAE(),
|
|
positive=positive,
|
|
negative=[],
|
|
sampler=object(),
|
|
sigmas=torch.tensor([1.0, 0.0]),
|
|
seed=70,
|
|
latent_template=latent_template,
|
|
num_frames_per_scene=5,
|
|
frame_overlap=2,
|
|
cfg=1.0,
|
|
max_scenes=8,
|
|
)
|
|
|
|
assert scene_count == 3
|
|
assert [call["seed"] for call in sample_calls] == [70, 71, 72]
|
|
assert [call["positive"][0][0].mean().item() for call in sample_calls] == [1.0, 2.0, 3.0]
|
|
assert "3 scene(s)" in status
|
|
assert "Scene 1: seed=70, text=Scene 1" in report
|
|
assert "Scene 3: seed=72, text=Scene 3" in report
|
|
|
|
|
|
def test_scene_chain_uses_scene_seed_metadata():
|
|
sample_calls.clear()
|
|
node = FunPackLTXAVSceneChainSampler()
|
|
latent_template = {"samples": torch.zeros(1, 2, 5, 3, 3)}
|
|
positive = [
|
|
(torch.ones(1, 2, 3), {"funpack_scene_text": "scene a", "funpack_scene_seed": 101}),
|
|
(torch.ones(1, 2, 3), {"funpack_scene_text": "scene b", "funpack_scene_seed": 202}),
|
|
(torch.ones(1, 2, 3), {"funpack_scene_text": "scene c", "funpack_scene_seed": 303}),
|
|
]
|
|
|
|
_, _, scene_count, report = node.sample(
|
|
model=object(),
|
|
vae=FakeVAE(),
|
|
positive=positive,
|
|
negative=[],
|
|
sampler=object(),
|
|
sigmas=torch.tensor([1.0, 0.0]),
|
|
seed=10,
|
|
latent_template=latent_template,
|
|
num_frames_per_scene=5,
|
|
frame_overlap=2,
|
|
cfg=1.0,
|
|
max_scenes=8,
|
|
)
|
|
|
|
assert scene_count == 3
|
|
assert [call["seed"] for call in sample_calls] == [101, 202, 303]
|
|
assert "seed=202" in report
|
|
|
|
|
|
def test_scene_chain_use_same_seed_forces_first_seed():
|
|
sample_calls.clear()
|
|
node = FunPackLTXAVSceneChainSampler()
|
|
latent_template = {"samples": torch.zeros(1, 2, 5, 3, 3)}
|
|
positive = [
|
|
(torch.ones(1, 2, 3), {"funpack_scene_text": "scene a", "funpack_scene_seed": 101}),
|
|
(torch.ones(1, 2, 3), {"funpack_scene_text": "scene b", "funpack_scene_seed": 202}),
|
|
(torch.ones(1, 2, 3), {"funpack_scene_text": "scene c"}),
|
|
]
|
|
|
|
_, _, scene_count, _ = node.sample(
|
|
model=object(),
|
|
vae=FakeVAE(),
|
|
positive=positive,
|
|
negative=[],
|
|
sampler=object(),
|
|
sigmas=torch.tensor([1.0, 0.0]),
|
|
seed=10,
|
|
latent_template=latent_template,
|
|
num_frames_per_scene=5,
|
|
frame_overlap=2,
|
|
cfg=1.0,
|
|
max_scenes=8,
|
|
use_same_seed=True,
|
|
)
|
|
|
|
assert scene_count == 3
|
|
assert [call["seed"] for call in sample_calls] == [101, 101, 101]
|
|
|
|
|
|
def test_scene_chain_preserves_nested_av_structure_and_audio_length():
|
|
sample_calls.clear()
|
|
node = FunPackLTXAVSceneChainSampler()
|
|
video = torch.zeros(1, 2, 5, 3, 3)
|
|
audio = torch.zeros(1, 1, 10, 4)
|
|
latent_template = {"samples": FakeNestedTensor([video, audio])}
|
|
positive = [scene_cond(0), scene_cond(1)]
|
|
|
|
latent, _, scene_count, _ = node.sample(
|
|
model=object(),
|
|
vae=FakeVAE(),
|
|
positive=positive,
|
|
negative=[],
|
|
sampler=object(),
|
|
sigmas=torch.tensor([1.0, 0.0]),
|
|
seed=20,
|
|
latent_template=latent_template,
|
|
num_frames_per_scene=5,
|
|
frame_overlap=2,
|
|
cfg=1.0,
|
|
max_scenes=8,
|
|
)
|
|
|
|
video_out, audio_out = latent["samples"].unbind()
|
|
assert scene_count == 2
|
|
assert video_out.shape[2] == 8
|
|
assert audio_out.shape[2] == 16
|
|
|
|
|
|
def test_scene_chain_default_max_is_eight_but_allows_more():
|
|
inputs = FunPackLTXAVSceneChainSampler.INPUT_TYPES()["required"]["max_scenes"][1]
|
|
assert inputs["default"] == 8
|
|
assert "max" not in inputs
|
|
|
|
sample_calls.clear()
|
|
node = FunPackLTXAVSceneChainSampler()
|
|
latent_template = {"samples": torch.zeros(1, 2, 3, 2, 2)}
|
|
positive = [scene_cond(index) for index in range(10)]
|
|
|
|
latent, status, scene_count, report = node.sample(
|
|
model=object(),
|
|
vae=FakeVAE(),
|
|
positive=positive,
|
|
negative=[],
|
|
sampler=object(),
|
|
sigmas=torch.tensor([1.0, 0.0]),
|
|
seed=30,
|
|
latent_template=latent_template,
|
|
num_frames_per_scene=3,
|
|
frame_overlap=0,
|
|
cfg=1.0,
|
|
max_scenes=10,
|
|
)
|
|
|
|
assert scene_count == 10
|
|
assert len(sample_calls) == 10
|
|
assert sample_calls[-1]["seed"] == 39
|
|
assert latent["samples"].shape[2] == 30
|
|
assert "10 scene(s)" in status
|
|
assert "Scene 10" in report
|
|
|
|
|
|
def test_scene_chain_can_append_i2v_template_as_hidden_guide():
|
|
sample_calls.clear()
|
|
node = FunPackLTXAVSceneChainSampler()
|
|
samples = torch.zeros(1, 2, 5, 1, 1)
|
|
samples[:, :, 0] = 7.0
|
|
mask = torch.ones(1, 1, 5, 1, 1)
|
|
mask[:, :, 0] = 0.0
|
|
latent_template = {"samples": samples, "noise_mask": mask}
|
|
positive = [scene_cond(0), scene_cond(1)]
|
|
|
|
latent, status, scene_count, _ = node.sample(
|
|
model=object(),
|
|
vae=FakeVAE(),
|
|
positive=positive,
|
|
negative=[],
|
|
sampler=object(),
|
|
sigmas=torch.tensor([1.0, 0.0]),
|
|
seed=40,
|
|
latent_template=latent_template,
|
|
num_frames_per_scene=5,
|
|
frame_overlap=2,
|
|
cfg=1.0,
|
|
max_scenes=2,
|
|
carry_i2v_guides=True,
|
|
)
|
|
|
|
second_call = sample_calls[1]
|
|
assert scene_count == 2
|
|
assert second_call["latent_image"].shape[2] == 6
|
|
assert torch.all(second_call["latent_image"][:, :, 2:5] == 0.0)
|
|
assert torch.all(second_call["latent_image"][:, :, 5] == 7.0)
|
|
assert torch.all(second_call["noise_mask"][:, :, :2] == 0.0)
|
|
assert torch.all(second_call["noise_mask"][:, :, 2:5] == 1.0)
|
|
assert torch.all(second_call["noise_mask"][:, :, 5] == 0.0)
|
|
assert "keyframe_idxs" in second_call["positive"][0][1]
|
|
assert second_call["positive"][0][1]["guide_attention_entries"][0]["pre_filter_count"] == 1
|
|
assert latent["samples"].shape[2] == 8
|
|
assert "i2v guide tokens=1 latent frame(s)" in status
|
|
|
|
|
|
def test_scene_chain_expands_compact_i2v_guide_mask_to_spatial_chunk_mask():
|
|
sample_calls.clear()
|
|
node = FunPackLTXAVSceneChainSampler()
|
|
samples = torch.zeros(1, 2, 5, 24, 3)
|
|
samples[:, :, 0] = 7.0
|
|
mask = torch.ones(1, 1, 5, 1, 1)
|
|
mask[:, :, 0] = 0.0
|
|
latent_template = {"samples": samples, "noise_mask": mask}
|
|
positive = [scene_cond(0), scene_cond(1)]
|
|
|
|
_, status, scene_count, _ = node.sample(
|
|
model=object(),
|
|
vae=FakeVAE(),
|
|
positive=positive,
|
|
negative=[],
|
|
sampler=object(),
|
|
sigmas=torch.tensor([1.0, 0.0]),
|
|
seed=45,
|
|
latent_template=latent_template,
|
|
num_frames_per_scene=5,
|
|
frame_overlap=2,
|
|
cfg=1.0,
|
|
max_scenes=2,
|
|
carry_i2v_guides=True,
|
|
)
|
|
|
|
second_call = sample_calls[1]
|
|
assert scene_count == 2
|
|
assert second_call["noise_mask"].shape == second_call["latent_image"].shape
|
|
assert torch.all(second_call["noise_mask"][:, :, 5] == 0.0)
|
|
assert "i2v guide tokens=1 latent frame(s)" in status
|
|
|
|
|
|
def test_scene_chain_does_not_carry_i2v_guides_by_default():
|
|
inputs = FunPackLTXAVSceneChainSampler.INPUT_TYPES()["required"]["carry_i2v_guides"][1]
|
|
assert inputs["default"] is False
|
|
|
|
sample_calls.clear()
|
|
node = FunPackLTXAVSceneChainSampler()
|
|
samples = torch.zeros(1, 2, 5, 1, 1)
|
|
samples[:, :, 0] = 7.0
|
|
mask = torch.ones(1, 1, 5, 1, 1)
|
|
mask[:, :, 0] = 0.0
|
|
latent_template = {"samples": samples, "noise_mask": mask}
|
|
positive = [scene_cond(0), scene_cond(1)]
|
|
|
|
_, status, scene_count, _ = node.sample(
|
|
model=object(),
|
|
vae=FakeVAE(),
|
|
positive=positive,
|
|
negative=[],
|
|
sampler=object(),
|
|
sigmas=torch.tensor([1.0, 0.0]),
|
|
seed=50,
|
|
latent_template=latent_template,
|
|
num_frames_per_scene=5,
|
|
frame_overlap=2,
|
|
cfg=1.0,
|
|
max_scenes=2,
|
|
)
|
|
|
|
second_call = sample_calls[1]
|
|
assert scene_count == 2
|
|
assert torch.all(second_call["latent_image"][:, :, 2] == 0.0)
|
|
assert torch.all(second_call["noise_mask"][:, :, :2] == 0.0)
|
|
assert torch.all(second_call["noise_mask"][:, :, 2:] == 1.0)
|
|
assert "i2v guide" not in status
|
|
|
|
|
|
def test_overlap_diagnostics_report_latent_blend_zone():
|
|
import json
|
|
|
|
node = FunPackLTXAVSceneChainSampler()
|
|
diag = node._build_overlap_diagnostics(
|
|
scene_count=2,
|
|
video_frames=97,
|
|
pixel_overlap=16,
|
|
latent_overlap=2,
|
|
time_scale=8,
|
|
transition_duration=16,
|
|
boundaries=[{
|
|
"between": [1, 2],
|
|
"boundary_latent": 12,
|
|
"pixel_frame": 89,
|
|
"effect": "crossfade",
|
|
}],
|
|
scene_runs=[
|
|
{"index": 1, "text": "hero walks", "encode_text": "hero walks", "mechanisms": []},
|
|
{"index": 2, "text": "hero runs", "encode_text": "hero runs", "mechanisms": ["latent_overlap(16px)"]},
|
|
],
|
|
carry_i2v_guides=True,
|
|
mid_scene_guide=False,
|
|
embed_guidance=True,
|
|
embed_guidance_strength=0.15,
|
|
embed_guidance_source="absolute",
|
|
)
|
|
assert diag["pixel_overlap"] == 16
|
|
blend = diag["boundaries"][0]["contamination_zones"]["latent_blend"]
|
|
assert blend["scene_prev_tail"] == [73, 88]
|
|
assert blend["scene_next_head"] == [89, 104]
|
|
assert diag["scenes"][0]["whole_scene_steering"] is True
|
|
assert any(g["mechanism"] == "embed_guidance" for g in diag["global_steering"])
|
|
# scene_boundaries output is JSON in the full sample() path
|
|
sample_calls.clear()
|
|
node = FunPackLTXAVSceneChainSampler()
|
|
positive = [scene_cond(0), scene_cond(1)]
|
|
latent_template = {"samples": torch.zeros(1, 2, 5, 1, 1)}
|
|
_, _, status, _, _, boundaries_json = node.sample(
|
|
model=object(),
|
|
vae=FakeVAE(),
|
|
positive=positive,
|
|
negative=[],
|
|
sampler=object(),
|
|
sigmas=torch.tensor([1.0, 0.0]),
|
|
seed=60,
|
|
latent_template=latent_template,
|
|
num_frames_per_scene=5,
|
|
frame_overlap=2,
|
|
cfg=1.0,
|
|
max_scenes=2,
|
|
embed_guidance=False,
|
|
)
|
|
parsed = json.loads(boundaries_json)
|
|
assert parsed["scene_count"] == 2
|
|
assert "boundaries" in parsed
|
|
assert "overlap_blend=2px" in status
|
|
|
|
|
|
def test_mixed_solo_applies_guides_on_first_chunk():
|
|
import json
|
|
|
|
sample_calls.clear()
|
|
node = FunPackLTXAVSceneChainSampler()
|
|
guides = json.dumps({
|
|
"stack_enabled": True,
|
|
"accumulate_prior": False,
|
|
"scenes": [[{"enabled": True, "source": "template", "frame_idx": 0, "apply_at": 0, "strength": 0.35}]],
|
|
})
|
|
latent_template = {
|
|
"samples": torch.zeros(1, 2, 5, 3, 3),
|
|
"noise_mask": torch.cat([torch.zeros(1, 1, 2, 1, 1), torch.ones(1, 1, 3, 1, 1)], dim=2),
|
|
}
|
|
positive = [scene_cond(0)]
|
|
|
|
latent, _images, status, scene_count, _report, boundaries_json = node.sample(
|
|
model=object(),
|
|
vae=FakeVAE(),
|
|
positive=positive,
|
|
negative=[],
|
|
sampler=object(),
|
|
sigmas=torch.tensor([1.0, 0.0]),
|
|
seed=90,
|
|
latent_template=latent_template,
|
|
num_frames_per_scene=5,
|
|
frame_overlap=0,
|
|
cfg=1.0,
|
|
max_scenes=1,
|
|
funpack_scene_guides=guides,
|
|
prompt={"n1": {"inputs": {}}},
|
|
unique_id="test-node",
|
|
)
|
|
|
|
assert scene_count == 1
|
|
chunk = sample_calls[0]["latent_image"]
|
|
frames = chunk["samples"].shape[2] if isinstance(chunk, dict) else chunk.shape[2]
|
|
assert frames == 7
|
|
parsed = json.loads(boundaries_json)
|
|
assert "custom_guide_stack" in parsed["scenes"][0]["mechanisms"]
|
|
|
|
|
|
def test_mixed_anchor_skips_frame_overlap():
|
|
import json
|
|
|
|
sample_calls.clear()
|
|
node = FunPackLTXAVSceneChainSampler()
|
|
node._load_image_tensor = lambda _fn: torch.ones(1, 8, 8, 3)
|
|
node._apply_img2video_to_video_latent = lambda _vae, _img, chunk, _strength: node._clone_latent(chunk)
|
|
|
|
positive = [scene_cond(0), scene_cond(1)]
|
|
latent_template = {"samples": torch.zeros(1, 2, 5, 3, 3)}
|
|
anchors = json.dumps({"1": {"filename": "anchor.png", "strength": 1.0}})
|
|
|
|
latent, _images, status, scene_count, _report, boundaries_json = node.sample(
|
|
model=object(),
|
|
vae=FakeVAE(),
|
|
positive=positive,
|
|
negative=[],
|
|
sampler=object(),
|
|
sigmas=torch.tensor([1.0, 0.0]),
|
|
seed=80,
|
|
latent_template=latent_template,
|
|
num_frames_per_scene=5,
|
|
frame_overlap=2,
|
|
cfg=1.0,
|
|
max_scenes=2,
|
|
funpack_scene_anchors=anchors,
|
|
prompt={"n1": {"inputs": {}}},
|
|
unique_id="test-node",
|
|
)
|
|
|
|
assert scene_count == 2
|
|
assert latent["samples"].shape[2] == 10
|
|
parsed = json.loads(boundaries_json)
|
|
mechs = parsed["scenes"][1]["mechanisms"]
|
|
assert "mixed_i2v_anchor" in mechs
|
|
assert not any("latent_overlap" in m for m in mechs)
|