136 lines
5.7 KiB
Python
136 lines
5.7 KiB
Python
"""First-step seed search: try a few seeds for one step, keep the one your taste likes.
|
|
|
|
Before the real run, each candidate seed gets exactly ONE denoising step, its
|
|
predicted picture is scored by a judge trained on how your rated clips looked
|
|
at their first step, and the run continues with the best seed. The first step
|
|
already carries most of what a rating reacts to (v4's probe: ~88% of the late
|
|
half's signal).
|
|
|
|
User verdict (v4): a win on simple prompts; no help on prompts that depend on
|
|
word trickery, which the first step can't read yet.
|
|
|
|
Costs one extra model pass per candidate. Candidate steps are marked as probes,
|
|
so no learning feature records a step the clip never took.
|
|
"""
|
|
|
|
import torch
|
|
from comfy.patcher_extension import WrappersMP
|
|
|
|
from ..._core import dit_hooks, log, registry, streams
|
|
|
|
ID = "explore_first_step"
|
|
TITLE = "First-step seed search"
|
|
MOUNT = "generation.sampling"
|
|
STAGE = "sampling"
|
|
CATEGORY = "sampling"
|
|
STATUS = "proven"
|
|
REQUIRES = ["temporal_latent"]
|
|
USES = ["taste_store"]
|
|
|
|
SETTINGS = {
|
|
"enabled": {
|
|
"type": "bool", "default": False,
|
|
"label": "Seed search",
|
|
"hint": "After 10+ rated clips (some liked, some disliked), tries a few seeds for one step each and keeps the one your taste likes.",
|
|
},
|
|
"candidates": {
|
|
"type": "int", "default": 3, "min": 2, "max": 8,
|
|
"label": "Seeds to try",
|
|
"hint": "Each costs one extra step. Your seed is always one of them.",
|
|
"when": {"enabled": True},
|
|
},
|
|
}
|
|
|
|
KIND = "first_step"
|
|
NAME = "first"
|
|
|
|
|
|
def _say(message):
|
|
log.once(f"{ID}:{message}", log.ALERT, "FunPack Seed search", message)
|
|
|
|
|
|
def install(patcher, values, key):
|
|
if not values.get("enabled"):
|
|
return None
|
|
taste = registry.current().ask("taste_store", patcher)
|
|
if taste is None:
|
|
_say("off: set a Taste key first -- it is what your ratings teach")
|
|
return None
|
|
count = int(values.get("candidates", 3))
|
|
live = {}
|
|
|
|
def fresh():
|
|
live["judge"] = taste.judge(KIND, NAME)
|
|
liked, disliked = taste.counts(KIND, NAME)
|
|
if live["judge"] is None:
|
|
log.once(f"{ID}:state", log.INFO, "FunPack Seed search",
|
|
f"key {taste.key!r}: learning ({liked} liked / {disliked} disliked; "
|
|
"searches from 10 rated clips with both kinds)")
|
|
|
|
captured = taste.collect(patcher, key, KIND, fresh=fresh)
|
|
|
|
# What the real run's first step looked like: what the judge learns from.
|
|
def apply_model(executor, x, t, *args, **kwargs):
|
|
out = executor(x, t, *args, **kwargs)
|
|
named = streams.model_args(args, kwargs)
|
|
to = named.get("transformer_options")
|
|
where = dit_hooks.current_step(to)
|
|
if dit_hooks.probing(to) or where is None or where[0] != 0:
|
|
return out
|
|
split = streams.video_of(out, named)
|
|
if split is not None and split[0].shape[0] == 1:
|
|
captured[NAME] = taste.describe(split[0].detach())
|
|
return out
|
|
|
|
patcher.add_wrapper_with_key(WrappersMP.APPLY_MODEL, key, apply_model)
|
|
|
|
def sampler_sample(executor, model_wrap, sigmas, extra_args, callback, noise,
|
|
latent_image=None, denoise_mask=None, disable_pbar=False):
|
|
run = lambda n: executor(model_wrap, sigmas, extra_args, callback, n,
|
|
latent_image, denoise_mask, disable_pbar)
|
|
judge = live.get("judge")
|
|
sampler = getattr(executor, "class_obj", None)
|
|
if judge is None or noise.shape[0] != 1 or not hasattr(sampler, "max_denoise"):
|
|
return run(noise)
|
|
best = pick(judge, model_wrap, sigmas, extra_args, noise, latent_image, denoise_mask,
|
|
sampler.max_denoise(model_wrap, sigmas))
|
|
return run(best)
|
|
|
|
def pick(judge, model_wrap, sigmas, extra_args, noise, latent_image, denoise_mask, max_denoise):
|
|
from comfy.samplers import KSamplerX0Inpaint
|
|
options = dict(extra_args.get("model_options") or {})
|
|
options["transformer_options"] = {**options.get("transformer_options", {}),
|
|
dit_hooks.PROBE: True}
|
|
named = {"latent_shapes": getattr(model_wrap.inner_model, "latent_shapes", None)}
|
|
seed = int(extra_args.get("seed") or 0)
|
|
sigma = sigmas[0] * noise.new_ones([noise.shape[0]])
|
|
scored = []
|
|
for k in range(count):
|
|
if k == 0:
|
|
candidate = noise # the seed you asked for
|
|
else:
|
|
gen = torch.Generator().manual_seed(seed + 7919 * k)
|
|
candidate = torch.randn(noise.shape, generator=gen).to(noise)
|
|
model_k = KSamplerX0Inpaint(model_wrap, sigmas)
|
|
model_k.latent_image, model_k.noise = latent_image, candidate
|
|
x = model_wrap.inner_model.model_sampling.noise_scaling(
|
|
sigmas[0], candidate, latent_image, max_denoise)
|
|
x0 = model_k(x, sigma, denoise_mask=denoise_mask, model_options=options,
|
|
seed=extra_args.get("seed"))
|
|
split = streams.video_of(x0, named)
|
|
if split is None:
|
|
_say("off this run: could not find the picture to score")
|
|
return noise
|
|
scored.append((judge.score(split[0]), k, candidate))
|
|
score, k, best = max(scored, key=lambda s: s[0])
|
|
log.info("FunPack Seed search",
|
|
("kept your seed" if k == 0 else f"switched to candidate {k + 1}")
|
|
+ f" of {count} (scores {', '.join(f'{s:.3f}' for s, _k, _c in scored)})")
|
|
return best
|
|
|
|
patcher.add_wrapper_with_key(WrappersMP.SAMPLER_SAMPLE, key, sampler_sample)
|
|
return f"{count} seeds, one step each; learns every run, searches from 10 rated clips"
|
|
|
|
|
|
PROVIDES = {"modifier": install}
|