digital-garbageandClaude Sonnet 5
99c784921d
Tries N candidate seeds with one throwaway denoising step each, scores each with the trained output-space value function, commits to the best-scoring seed for the real full pass. Candidates are scored fully isolated from the scene's own guidance stack (model_function_wrapper chain, context_windows, v2a_grad_scale, identity_overlap -- each torn down and reinstalled around the candidate-scoring window) against the scene's fully-finalized conditioning. Any failure in the value function, schedule, or the isolation cycle itself degrades to a silent no-op, declared in the run report, never a crash. Also hardens _install_context_windows/_install_v2a_scale to roll back partial state on failure (a leaked context_handler key was silently re-enabling windowing for later scenes) and fixes the original non-explore v2a install call site to degrade like its sibling mechanisms instead of aborting the render. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
238 lines
10 KiB
Python
238 lines
10 KiB
Python
"""_select_best_seed: branch-then-commit seed search at the first step only.
|
|
|
|
See samplers.py's docstring on _select_best_seed for the mechanism. This isolates it from
|
|
the full sample() harness in test_scene_chain_sampler.py deliberately: that file's
|
|
fake_sample_custom never touches model.model_options["model_function_wrapper"] at all (it
|
|
computes its stand-in result directly from the inputs), so a full-sample() test would never
|
|
actually exercise the observer-wrapper capture this method depends on -- it would silently
|
|
pass while testing nothing. The fakes here call the installed wrapper directly, the way a
|
|
real one-step comfy.sample.sample_custom call would.
|
|
"""
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
import _comfy_stubs
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
_comfy_stubs.install_module("comfy.sample", prepare_noise=lambda *a, **k: None,
|
|
sample_custom=lambda *a, **k: None)
|
|
|
|
from samplers import FunPackLTXAVSceneChainSampler # noqa: E402
|
|
|
|
SEED_OFFSET = 999983 # must match the offset _select_best_seed derives candidate seeds with
|
|
|
|
|
|
class _FakeValueFn:
|
|
"""Ranks a captured tensor by its own mean -- deterministic, easy to reason about,
|
|
standing in for the real trained MLP."""
|
|
def is_ready(self):
|
|
return True
|
|
|
|
def compress(self, x):
|
|
return x
|
|
|
|
def forward(self, x):
|
|
return x.mean()
|
|
|
|
|
|
class _NotReadyValueFn(_FakeValueFn):
|
|
def is_ready(self):
|
|
return False
|
|
|
|
|
|
@pytest.fixture
|
|
def refiner():
|
|
return FunPackLTXAVSceneChainSampler.__new__(FunPackLTXAVSceneChainSampler)
|
|
|
|
|
|
def _fake_model():
|
|
return types.SimpleNamespace(model_options={})
|
|
|
|
|
|
def _install_fakes(monkeypatch, denoised_for_seed, prepare_noise_calls=None, sample_calls=None):
|
|
def fake_prepare_noise(samples, seed, noise_inds=None):
|
|
if prepare_noise_calls is not None:
|
|
prepare_noise_calls.append(seed)
|
|
return torch.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):
|
|
if sample_calls is not None:
|
|
sample_calls.append(seed)
|
|
wrapper = model.model_options.get("model_function_wrapper")
|
|
denoised = denoised_for_seed(seed, latent_image)
|
|
if wrapper is not None:
|
|
wrapper(lambda x, t, **c: denoised, {"input": noise, "timestep": torch.tensor([1.0])})
|
|
return latent_image
|
|
|
|
monkeypatch.setattr(sys.modules["comfy.sample"], "prepare_noise", fake_prepare_noise, raising=False)
|
|
monkeypatch.setattr(sys.modules["comfy.sample"], "sample_custom", fake_sample_custom, raising=False)
|
|
|
|
|
|
def test_the_highest_scoring_candidate_seed_wins(monkeypatch, refiner):
|
|
model = _fake_model()
|
|
sample_calls = []
|
|
# the fake model's "prediction" is literally the seed's own value broadcast into the
|
|
# latent -- the fake value_fn ranks by mean, so the LARGEST seed must always win.
|
|
_install_fakes(monkeypatch, lambda seed, latent: torch.full_like(latent, float(seed)),
|
|
sample_calls=sample_calls)
|
|
|
|
latent = {"samples": torch.zeros(1, 4, 2, 2, 2)}
|
|
sigmas = torch.tensor([1.0, 0.5, 0.0])
|
|
winner = refiner._select_best_seed(
|
|
model, object(), sigmas, seed=100, cfg=1.0, positive=[], negative=[],
|
|
latent=latent, n_candidates=3, value_fn=_FakeValueFn())
|
|
|
|
assert winner == 100 + 2 * SEED_OFFSET
|
|
assert sample_calls == [100, 100 + SEED_OFFSET, 100 + 2 * SEED_OFFSET]
|
|
# the observer wrapper must be stripped after every candidate, success or failure --
|
|
# the REAL run that follows must never inherit it.
|
|
assert model.model_options.get("model_function_wrapper") is None
|
|
|
|
|
|
def test_a_lower_seed_can_win_when_it_scores_higher(monkeypatch, refiner):
|
|
"""Not just 'largest seed wins' by construction -- the WINNING seed is whichever one the
|
|
scorer actually preferred, which here is deliberately the FIRST (lowest) candidate."""
|
|
model = _fake_model()
|
|
|
|
def denoised_for_seed(seed, latent):
|
|
# seed 100 (candidate 0) scores highest; the other two score lower.
|
|
value = 10.0 if seed == 100 else 1.0
|
|
return torch.full_like(latent, value)
|
|
|
|
_install_fakes(monkeypatch, denoised_for_seed)
|
|
latent = {"samples": torch.zeros(1, 4, 2, 2, 2)}
|
|
sigmas = torch.tensor([1.0, 0.5, 0.0])
|
|
winner = refiner._select_best_seed(
|
|
model, object(), sigmas, seed=100, cfg=1.0, positive=[], negative=[],
|
|
latent=latent, n_candidates=3, value_fn=_FakeValueFn())
|
|
assert winner == 100
|
|
|
|
|
|
def test_a_value_function_that_is_not_ready_is_a_pure_no_op(monkeypatch, refiner):
|
|
model = _fake_model()
|
|
calls = []
|
|
_install_fakes(monkeypatch, lambda seed, latent: torch.full_like(latent, float(seed)),
|
|
sample_calls=calls)
|
|
latent = {"samples": torch.zeros(1, 4, 2, 2, 2)}
|
|
sigmas = torch.tensor([1.0, 0.5, 0.0])
|
|
winner = refiner._select_best_seed(
|
|
model, object(), sigmas, seed=100, cfg=1.0, positive=[], negative=[],
|
|
latent=latent, n_candidates=3, value_fn=_NotReadyValueFn())
|
|
assert winner == 100
|
|
assert calls == [] # no candidate work was ever attempted
|
|
|
|
|
|
def test_none_value_function_is_a_pure_no_op(refiner):
|
|
latent = {"samples": torch.zeros(1, 4, 2, 2, 2)}
|
|
sigmas = torch.tensor([1.0, 0.5, 0.0])
|
|
winner = refiner._select_best_seed(
|
|
_fake_model(), object(), sigmas, seed=100, cfg=1.0, positive=[], negative=[],
|
|
latent=latent, n_candidates=3, value_fn=None)
|
|
assert winner == 100
|
|
|
|
|
|
def test_a_schedule_too_short_to_slice_one_step_is_a_no_op(refiner):
|
|
latent = {"samples": torch.zeros(1, 4, 2, 2, 2)}
|
|
winner = refiner._select_best_seed(
|
|
_fake_model(), object(), torch.tensor([1.0]), seed=100, cfg=1.0, positive=[],
|
|
negative=[], latent=latent, n_candidates=3, value_fn=_FakeValueFn())
|
|
assert winner == 100
|
|
|
|
|
|
def test_fewer_than_two_candidates_is_a_no_op(refiner):
|
|
latent = {"samples": torch.zeros(1, 4, 2, 2, 2)}
|
|
sigmas = torch.tensor([1.0, 0.5, 0.0])
|
|
winner = refiner._select_best_seed(
|
|
_fake_model(), object(), sigmas, seed=100, cfg=1.0, positive=[], negative=[],
|
|
latent=latent, n_candidates=1, value_fn=_FakeValueFn())
|
|
assert winner == 100
|
|
|
|
|
|
def test_a_candidate_whose_sample_custom_call_raises_is_skipped_not_fatal(monkeypatch, refiner):
|
|
model = _fake_model()
|
|
|
|
def fake_prepare_noise(samples, seed, noise_inds=None):
|
|
return torch.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):
|
|
if seed == 100:
|
|
raise RuntimeError("candidate 0 boom")
|
|
wrapper = model.model_options.get("model_function_wrapper")
|
|
denoised = torch.full_like(latent_image, float(seed))
|
|
if wrapper is not None:
|
|
wrapper(lambda x, t, **c: denoised, {"input": noise, "timestep": torch.tensor([1.0])})
|
|
return latent_image
|
|
|
|
monkeypatch.setattr(sys.modules["comfy.sample"], "prepare_noise", fake_prepare_noise, raising=False)
|
|
monkeypatch.setattr(sys.modules["comfy.sample"], "sample_custom", fake_sample_custom, raising=False)
|
|
|
|
latent = {"samples": torch.zeros(1, 4, 2, 2, 2)}
|
|
sigmas = torch.tensor([1.0, 0.5, 0.0])
|
|
winner = refiner._select_best_seed(
|
|
model, object(), sigmas, seed=100, cfg=1.0, positive=[], negative=[],
|
|
latent=latent, n_candidates=3, value_fn=_FakeValueFn())
|
|
# candidate 0 (seed=100) raised and is excluded from scoring entirely -- the winner
|
|
# must be the better of the two SURVIVING candidates, not the seed that crashed, and
|
|
# the crash must not propagate (the real run must still get to happen).
|
|
assert winner == 100 + 2 * SEED_OFFSET
|
|
# the wrapper is still restored to None even after the raising candidate.
|
|
assert model.model_options.get("model_function_wrapper") is None
|
|
|
|
|
|
def test_every_candidate_failing_falls_back_to_the_original_seed(monkeypatch, refiner):
|
|
model = _fake_model()
|
|
|
|
def fake_prepare_noise(samples, seed, noise_inds=None):
|
|
return torch.zeros_like(samples)
|
|
|
|
def fake_sample_custom(*a, **k):
|
|
raise RuntimeError("every candidate boom")
|
|
|
|
monkeypatch.setattr(sys.modules["comfy.sample"], "prepare_noise", fake_prepare_noise, raising=False)
|
|
monkeypatch.setattr(sys.modules["comfy.sample"], "sample_custom", fake_sample_custom, raising=False)
|
|
|
|
latent = {"samples": torch.zeros(1, 4, 2, 2, 2)}
|
|
sigmas = torch.tensor([1.0, 0.5, 0.0])
|
|
winner = refiner._select_best_seed(
|
|
model, object(), sigmas, seed=100, cfg=1.0, positive=[], negative=[],
|
|
latent=latent, n_candidates=3, value_fn=_FakeValueFn())
|
|
assert winner == 100
|
|
assert model.model_options.get("model_function_wrapper") is None
|
|
|
|
|
|
def test_only_a_two_element_probe_schedule_is_used_regardless_of_the_real_schedule_length(monkeypatch, refiner):
|
|
"""The real run's full sigma schedule (dozens of steps) must never be handed to the
|
|
throwaway candidate calls -- only a 2-element slice (exactly one step) may reach
|
|
sample_custom, which is the entire cost guarantee this feature makes."""
|
|
model = _fake_model()
|
|
seen_lengths = []
|
|
|
|
def fake_prepare_noise(samples, seed, noise_inds=None):
|
|
return torch.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):
|
|
seen_lengths.append(int(sigmas.numel()))
|
|
wrapper = model.model_options.get("model_function_wrapper")
|
|
denoised = torch.full_like(latent_image, float(seed))
|
|
if wrapper is not None:
|
|
wrapper(lambda x, t, **c: denoised, {"input": noise, "timestep": torch.tensor([1.0])})
|
|
return latent_image
|
|
|
|
monkeypatch.setattr(sys.modules["comfy.sample"], "prepare_noise", fake_prepare_noise, raising=False)
|
|
monkeypatch.setattr(sys.modules["comfy.sample"], "sample_custom", fake_sample_custom, raising=False)
|
|
|
|
latent = {"samples": torch.zeros(1, 4, 2, 2, 2)}
|
|
sigmas = torch.linspace(1.0, 0.0, steps=40) # a realistic, long real-run schedule
|
|
refiner._select_best_seed(
|
|
model, object(), sigmas, seed=100, cfg=1.0, positive=[], negative=[],
|
|
latent=latent, n_candidates=3, value_fn=_FakeValueFn())
|
|
assert seen_lengths == [2, 2, 2]
|