diff --git a/conditioning.py b/conditioning.py index 282f9303..9ec4cf55 100644 --- a/conditioning.py +++ b/conditioning.py @@ -13799,14 +13799,16 @@ class FunPackStudio: # out of the sampler through a denoiser proxy — the same trick ALG uses. # Wrapping keeps the sampler you chose; it does not replace it. _sharp = float(cfg.get("ksampler_sharpness", 0.0) or 0.0) - if _sharp > 0.0: + _spatial = float(cfg.get("ksampler_latent_sharpness", 0.0) or 0.0) + if _sharp > 0.0 or _spatial > 0.0: try: from .samplers import _sharpen_wrap_sampler except ImportError: from samplers import _sharpen_wrap_sampler _wrapped = _sharpen_wrap_sampler( sampler, _sharp, - float(cfg.get("ksampler_sharpen_start_pct", 0.35) or 0.0)) + float(cfg.get("ksampler_sharpen_start_pct", 0.35) or 0.0), + spatial=_spatial) if _wrapped is not None: sampler = _wrapped else: diff --git a/movie_editor/frontend/sampler_panel.js b/movie_editor/frontend/sampler_panel.js index 2f3bff26..72592901 100644 --- a/movie_editor/frontend/sampler_panel.js +++ b/movie_editor/frontend/sampler_panel.js @@ -84,7 +84,7 @@ type: type || "Hybrid Euler 2S", sigmas: "", hybrid: defaultHybrid(), distilled: defaultDistilled(), ksampler_name: "euler", ksampler_steps: 8, ksampler_scheduler: KSAMPLER_USER_SIGMAS, - ksampler_sharpness: 0.0, ksampler_sharpen_start_pct: 0.35, + ksampler_sharpness: 0.0, ksampler_latent_sharpness: 0.0, ksampler_sharpen_start_pct: 0.35, }; } @@ -378,7 +378,13 @@ hint(container, "The same free unsharp mask FunPack's own samplers have, driven from " + "outside this sampler's loop. Recovers the fine detail that goes missing " + "at low resolution. 0 = off, 0.2-0.4 typical."); - if (+cfg.ksampler_sharpness > 0) { + row(container, "latent sharpness", + numCtrl(cfg.ksampler_latent_sharpness, 0, 2, 0.05, dk + "-kls", + (v) => { cfg.ksampler_latent_sharpness = v; saveNow(); })); + hint(container, "Sharpens the picture's edges in latent space on the late steps (unsharp " + + "mask on each frame's H/W); the model sees it re-noised next step and can " + + "agree or undo. Audio untouched. 0 = off, try 0.3-0.6. Experimental."); + if (+cfg.ksampler_sharpness > 0 || +cfg.ksampler_latent_sharpness > 0) { row(container, "sharpen last %", numCtrl(cfg.ksampler_sharpen_start_pct, 0, 1, 0.05, dk + "-kqsp", (v) => { cfg.ksampler_sharpen_start_pct = v; save(); })); diff --git a/samplers.py b/samplers.py index 5b08c3c0..e1f9b4e3 100644 --- a/samplers.py +++ b/samplers.py @@ -590,6 +590,36 @@ def _rescue_denoised(denoised, x, sigma, refinement_key, target, threshold, stre return denoised +def _apply_spatial_sharpness(denoised, model, amount): + """Spatial unsharp on the x0 prediction's VIDEO stream: x + amount * (x - blur3x3(x)). + + `quality_sharpness` sharpens along TIME (this step's guess against the last step's); + this one sharpens along H/W, the picture's actual edges. The model sees the sharpened + guess re-noised on the next step, so it can agree with or undo it - which is what keeps + it from reading as a post filter. Audio is untouched by construction: the packed latent + is sliced to the video span and only its H/W axes are filtered. No-op (unchanged + tensor) when the span cannot be verified. + """ + if not amount or amount <= 0.0: + return denoised + span = _video_span(model, denoised) + if span is None or len(span[2]) != 5: + return denoised + off, size, shape = span + b, c, t, h, w = (int(v) for v in shape) + if h < 3 or w < 3: + return denoised + flat = denoised[..., off:off + size].reshape(b, c, t, h, w) + frames = flat.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) + blurred = torch.nn.functional.avg_pool2d( + torch.nn.functional.pad(frames.float(), (1, 1, 1, 1), mode="replicate"), 3, stride=1) + sharp = frames.float() + float(amount) * (frames.float() - blurred) + out = denoised.clone() + out[..., off:off + size] = (sharp.to(denoised.dtype).reshape(b, t, c, h, w) + .permute(0, 2, 1, 3, 4).reshape(*denoised.shape[:-1], size)) + return out + + def _apply_quality_sharpness(denoised, prev_denoised, sharpness): """E: temporal-average unsharp on the x0 prediction during the quality phase. @@ -1234,11 +1264,12 @@ class _SharpenDenoiser: this composes with the ALG proxy in either order. """ - _OWN = ("_inner", "_amount", "_thr", "_mask", "_prev", "_mask_done") + _OWN = ("_inner", "_amount", "_spatial", "_thr", "_mask", "_prev", "_mask_done") - def __init__(self, inner, amount, threshold, mask): + def __init__(self, inner, amount, threshold, mask, spatial=0.0): self._inner = inner self._amount = float(amount) + self._spatial = float(spatial or 0.0) self._thr = float(threshold) self._mask = mask self._prev = None @@ -1276,9 +1307,12 @@ class _SharpenDenoiser: self._prev = None return denoised try: - denoised = _video_only( - _apply_quality_sharpness(denoised, self._prev, self._amount), - denoised, self._mask) + if self._amount > 0.0: + denoised = _video_only( + _apply_quality_sharpness(denoised, self._prev, self._amount), + denoised, self._mask) + if self._spatial > 0.0: + denoised = _apply_spatial_sharpness(denoised, self._inner, self._spatial) except Exception as _e: # noqa: BLE001 _log.failed("FunPackStudio", "quality sharpness", _e, "this evaluation keeps the unsharpened prediction") @@ -1295,14 +1329,16 @@ class _SharpenDenoiser: self._mask = None -def _sharpen_wrap_sampler(sampler, sharpness, start_pct): - """A SAMPLER equivalent to `sampler` with quality sharpness applied around its model. +def _sharpen_wrap_sampler(sampler, sharpness, start_pct, spatial=0.0): + """A SAMPLER equivalent to `sampler` with quality sharpness (temporal) and/or latent + sharpness (spatial) applied around its model, over the last `start_pct` of the schedule. - Returns `sampler` unchanged when sharpness is off, and None when the sampler exposes no + Returns `sampler` unchanged when both are off, and None when the sampler exposes no `sampler_function` to wrap. """ sharpness = max(0.0, min(1.0, float(sharpness or 0.0))) - if sharpness <= 0.0: + spatial = max(0.0, min(2.0, float(spatial or 0.0))) + if sharpness <= 0.0 and spatial <= 0.0: return sampler fn = getattr(sampler, "sampler_function", None) if fn is None: @@ -1322,12 +1358,15 @@ def _sharpen_wrap_sampler(sampler, sharpness, start_pct): return fn(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, **options) mask = _packed_video_mask(model, x) - print(f"[FunPack] quality sharpness {sharpness:.2f} on the last " + what = " + ".join(p for p in ( + f"quality sharpness {sharpness:.2f}" if sharpness > 0 else "", + f"latent sharpness {spatial:.2f}" if spatial > 0 else "") if p) + print(f"[FunPack] {what} on the last " f"{start_pct * 100:.0f}% of the schedule (sigma <= {thr:.4f}) — driven from " f"outside {getattr(fn, '__name__', 'the sampler')}'s loop, off the sigma of " f"each model call" f"{'' if mask is not None else '; single-stream latent, no audio to protect'}") - proxy = _SharpenDenoiser(model, sharpness, thr, mask) + proxy = _SharpenDenoiser(model, sharpness, thr, mask, spatial=spatial) try: return fn(proxy, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, **options) diff --git a/tests/test_quality_sharpness_ksampler.py b/tests/test_quality_sharpness_ksampler.py index bfcf0226..2fdadfab 100644 --- a/tests/test_quality_sharpness_ksampler.py +++ b/tests/test_quality_sharpness_ksampler.py @@ -208,3 +208,40 @@ def test_studio_keeps_sampling_when_the_sampler_cannot_be_wrapped(monkeypatch, c ) assert sampler == "SAMPLER::euler" assert "no sampler_function to wrap" in capsys.readouterr().out + + +# --- latent (spatial) sharpness -------------------------------------------------------- + + +def _packed_model(video_shape, audio_shape): + """A denoiser carrying comfy's packed-latent stream shapes where _video_span looks.""" + ls = types.SimpleNamespace(cond=[video_shape, audio_shape]) + conds = {"positive": [{"model_conds": {"latent_shapes": ls}}]} + return types.SimpleNamespace(inner_model=types.SimpleNamespace(conds=conds)) + + +def test_latent_sharpness_is_an_unsharp_mask_on_video_frames_and_leaves_audio_alone(): + torch.manual_seed(0) + vshape, ashape = (1, 2, 3, 4, 5), (1, 3, 2, 7) + video = torch.randn(vshape) + audio = torch.randn(ashape) + packed = torch.cat([video.reshape(1, 1, -1), audio.reshape(1, 1, -1)], dim=-1) + out = samplers._apply_spatial_sharpness(packed, _packed_model(vshape, ashape), 0.5) + n_v = video.numel() + assert torch.equal(out[..., n_v:], packed[..., n_v:]) # audio bit-exact + frames = video.permute(0, 2, 1, 3, 4).reshape(-1, 2, 4, 5) + blurred = torch.nn.functional.avg_pool2d( + torch.nn.functional.pad(frames, (1, 1, 1, 1), mode="replicate"), 3, stride=1) + want = (frames + 0.5 * (frames - blurred)).reshape(1, 3, 2, 4, 5).permute(0, 2, 1, 3, 4) + assert torch.allclose(out[..., :n_v].reshape(vshape), want, atol=1e-6) + + +def test_latent_sharpness_without_a_verifiable_span_changes_nothing(): + x = torch.randn(1, 1, 30) + assert torch.equal(samplers._apply_spatial_sharpness(x, _Model([x]), 0.5), x) + + +def test_latent_sharpness_alone_still_wraps_and_temporal_alone_still_works(): + fn = lambda m, x, sig, **k: x # noqa: E731 + assert samplers._sharpen_wrap_sampler(_sampler(fn), 0.0, 0.5, spatial=0.4) is not _sampler(fn) + assert samplers._sharpen_wrap_sampler(_sampler(fn), 0.0, 0.5, spatial=0.0).sampler_function is fn