feat: latent sharpness -- spatial unsharp on the late x0 predictions (KSampler branch)

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

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:

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(); }));

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)

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