Two stages: optional latent upscaler for the hand-over (bicubic stays the default); upscaler file loaded once and kept; credit upscaler authors

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Tw9fjVWCMpMsh4XkZmnKnh

modules/sampling/two_stage/__init__.py

@@ -12,6 +12,8 @@ It installs on the model, so it works with any sampler. Not yet run on a real H3
 
 from comfy.patcher_extension import WrappersMP
 
+import folder_paths
+
 from ..._core import log
 
 from .stages import keep_whole_schedule, make_wrapper
@@ -24,6 +26,18 @@ CATEGORY = "sampling"
 STATUS = "experimental"
 REQUIRES = ["audio_stream", "adaln_modalities"]
 
+
+def _files():
+    try:
+        return folder_paths.get_filename_list("latent_upscale_models")
+    except Exception:          # noqa: BLE001 -- folder not registered is the same as no files
+        return []
+
+
+def _options():
+    return [{"value": f, "label": f} for f in _files()] or [{"value": "", "label": "No file in models/latent_upscale_models"}]
+
+
 SETTINGS = {
     "enabled": {
         "type": "bool", "default": False,
@@ -43,6 +57,22 @@ SETTINGS = {
         "hint": "Out of the run's own steps. At least one step is always left for the full size.",
         "when": {"enabled": True},
     },
+    "resize": {
+        "type": "enum", "default": "bicubic",
+        "options": [{"value": "bicubic", "label": "Bicubic"}, {"value": "upscaler", "label": "Latent upscaler"}],
+        "label": "How the hand-over is enlarged",
+        "hint": "Bicubic is plain and free. Latent upscaler runs H3's learned latent upscaler (models/latent_upscale_models) "
+                "on the clean guess, so the added detail is learned instead of blurred; it costs one small network pass "
+                "once per run. The first stage size is kept, the upscaler is run at the ratio that reaches the full size.",
+        "when": {"enabled": True},
+    },
+    "upscaler": {
+        "type": "enum", "default": (_files() or [""])[0],
+        "options": _options(),
+        "label": "Latent upscaler file",
+        "hint": "The H3 latent upscaler file. The list is read again each time Settings opens.",
+        "when": {"enabled": True, "resize": "upscaler"},
+    },
 }
 
 
@@ -50,9 +80,10 @@ def install(patcher, values, key):
     if not values.get("enabled"):
         return None
     size, steps = float(values.get("size", 0.7)), int(values.get("steps", 2))
-    patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, key, make_wrapper(size, steps))
+    resize = values.get("upscaler") if values.get("resize") == "upscaler" and values.get("upscaler") else None
+    patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, key, make_wrapper(size, steps, resize))
     patcher.add_wrapper_with_key(WrappersMP.SAMPLER_SAMPLE, f"{key}.schedule", keep_whole_schedule)
-    return f"first {steps} step(s) at {size:.0%} of the size"
+    return f"first {steps} step(s) at {size:.0%} of the size" + (f", enlarged by {resize}" if resize else "")
 
 
 PROVIDES = {"modifier": install}

modules/sampling/two_stage/stages.py

@@ -51,24 +51,51 @@ def keep_whole_schedule(executor, guider, sigmas, extra_args, *args, **kwargs):
     return executor(guider, sigmas, extra_args, *args, **kwargs)
 
 
-def make_wrapper(size, steps):
+_CACHE = {}
+
+
+def _upscaler(name):
+    """Built once per file (and per change of the file on disk); the network sits on the offload device between runs."""
+    import os
+    import folder_paths
+    from ...latent.second_pass.nodes import FOLDER, load_upscaler
+    path = folder_paths.get_full_path_or_raise(FOLDER, name)
+    stamp = (path, os.path.getmtime(path))
+    if _CACHE.get("stamp") != stamp:
+        _CACHE.clear()
+        _CACHE.update(stamp=stamp, model=load_upscaler(name))
+    return _CACHE["model"]
+
+
+def enlarge(video, h, w, upscaler=None):
+    """The hand-over: the first stage's clean guess at the full grid. With an upscaler, its learned 2D resize at the
+    ratio that gets nearest (it works in even steps), then a bicubic trim to the exact grid; without, plain bicubic."""
+    from ...latent.second_pass.nodes import downscale_to, run_upscaler
+    if upscaler is None:
+        return downscale_to(video, h, w)
+    ratio = (h / video.shape[-2] + w / video.shape[-1]) / 2.0
+    if ratio <= 1.0:
+        return video
+    return downscale_to(run_upscaler(upscaler, video, ratio), h, w)
+
+
+def make_wrapper(size, steps, resize=None):
     def outer(executor, *args, **kwargs):
         # Every OUTER_SAMPLE wrapper behind this one belongs around the WHOLE run, not around each stage (Noise
         # merging's per-run state, Taste, STAS' report ...). They are moved to wrap the staged run as one call.
         from comfy.patcher_extension import WrapperExecutor
         rest = executor.wrappers[executor.idx + 1:]
-        staged = lambda *a, **k: _staged(executor.original, executor.class_obj, size, steps, *a, **k)
+        staged = lambda *a, **k: _staged(executor.original, executor.class_obj, size, steps, resize, *a, **k)
         return WrapperExecutor.new_class_executor(staged, executor.class_obj, rest, 0).execute(*args, **kwargs)
 
     return outer
 
 
-def _staged(original, guider, size, steps, *args, **kwargs):
+def _staged(original, guider, size, steps, resize, *args, **kwargs):
     run = original
     import comfy.nested_tensor
     import comfy.utils
     import latent_preview
-    from ...latent.second_pass.nodes import downscale_to
 
     noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed = (list(args) + [None] * 8)[:8]
     latent_shapes = kwargs.get("latent_shapes")
@@ -105,6 +132,12 @@ def _staged(original, guider, size, steps, *args, **kwargs):
             nested = lambda t: comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(t, shapes1))
             preview(step, nested(x0), nested(x), total)
 
+    upscaler = None
+    if resize:
+        try:
+            upscaler = _upscaler(resize)
+        except Exception as e:                      # noqa: BLE001 -- a missing file must not lose the run
+            log.warning(TITLE, f"latent upscaler not used, bicubic instead: {e}")
     log.info(TITLE, f"first {steps} step(s) at latent {h1}x{w1} instead of {h}x{w}, then {len(second) - 1} at the full size")
     setattr(guider, FULL, sigmas)
     setattr(guider, dit_hooks.STAGE, 1)
@@ -122,7 +155,7 @@ def _staged(original, guider, size, steps, *args, **kwargs):
     done = guider.model_patcher.model.process_latent_out(comfy.nested_tensor.NestedTensor((video, audio)).cpu())
     video, audio = done.unbind()
     handed, _ = comfy.utils.pack_latents(
-        (downscale_to(video, h, w).to(latent_image.device, latent_image.dtype), audio.to(latent_image.device, latent_image.dtype)))
+        (enlarge(video, h, w, upscaler).to(latent_image.device, latent_image.dtype), audio.to(latent_image.device, latent_image.dtype)))
 
     guider.conds = raw_conds                        # the first stage encoded the originals for its own size
     shifted = None

modules/sampling/two_stage/tests/test_two_stage.py

@@ -90,3 +90,55 @@ def test_the_whole_schedule_lives_where_the_stage_schedule_does():
     keep_whole_schedule(lambda *a, **k: None, guider, stage, extra)
     got = extra["model_options"]["transformer_options"]["sample_sigmas"]
     assert got.device == stage.device and got.dtype == stage.dtype
+
+
+@pytest.fixture
+def upscaler_file(tmp_path, monkeypatch):
+    import folder_paths
+    from safetensors.torch import save_file
+    from modules.models.minimax_h3.latent_upscaler import H3LatentResizer3D
+    torch.manual_seed(0)
+    net = H3LatentResizer3D(channels=32, in_blocks=1, out_blocks=1)
+    save_file(net.state_dict(), str(tmp_path / "h3_up.safetensors"))
+    monkeypatch.setitem(folder_paths.folder_names_and_paths, "latent_upscale_models",
+                        ([str(tmp_path)], {".safetensors"}))
+    folder_paths.filename_list_cache.pop("latent_upscale_models", None)
+    from modules.sampling.two_stage import stages
+    stages._CACHE.clear()
+    return "h3_up.safetensors"
+
+
+def test_the_latent_upscaler_hands_over_the_full_grid(model, upscaler_file):
+    from modules.sampling.two_stage import stages
+    note, out, calls = _run(model, {"enabled": True, "size": 0.5, "steps": 2,
+                                    "resize": "upscaler", "upscaler": upscaler_file})
+    assert "enlarged by h3_up.safetensors" in note
+    assert calls == [(4, 4), (4, 4), (8, 8), (8, 8)]
+    video, _audio = out.unbind()
+    assert video.shape == (1, 24, 2, 8, 8) and torch.isfinite(video).all()
+    first = stages._CACHE["model"]
+    _run(model, {"enabled": True, "size": 0.5, "steps": 2, "resize": "upscaler", "upscaler": upscaler_file})
+    assert stages._CACHE["model"] is first                  # loaded once, kept
+
+
+def test_a_ratio_that_does_not_fit_is_trimmed_to_the_exact_grid(upscaler_file):
+    from modules.sampling.two_stage import stages
+    net = stages._upscaler(upscaler_file)
+    video = torch.randn(1, 24, 2, 6, 10)
+    out = stages.enlarge(video, 9, 15, net)                 # 1.5x: the net works in even sizes
+    assert out.shape == (1, 24, 2, 9, 15) and out.dtype == video.dtype
+    assert stages.enlarge(video, 6, 10, net) is video
+
+
+def test_a_missing_upscaler_falls_back_to_bicubic(model, upscaler_file):
+    note, out, calls = _run(model, {"enabled": True, "size": 0.5, "steps": 2,
+                                    "resize": "upscaler", "upscaler": "gone.safetensors"})
+    assert calls == [(4, 4), (4, 4), (8, 8), (8, 8)] and torch.isfinite(out.unbind()[0]).all()
+
+
+def test_the_upscaler_follows_the_latents_dtype(upscaler_file):
+    from modules.sampling.two_stage import stages
+    net = stages._upscaler(upscaler_file)
+    video = torch.randn(1, 24, 2, 4, 4, device="cpu", dtype=torch.bfloat16)
+    out = stages.enlarge(video, 8, 8, net)
+    assert out.dtype == torch.bfloat16 and out.device == video.device