From f66df804fecdbba1f0506f14d458fa07db2f30e6 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 13:49:44 +0000 Subject: [PATCH] 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 Claude-Session: https://claude.ai/code/session_01Tw9fjVWCMpMsh4XkZmnKnh --- CHANGELOG.md | 1 + README.md | 1 + modules/latent/second_pass/nodes.py | 31 ++++++----- modules/sampling/two_stage/__init__.py | 35 ++++++++++++- modules/sampling/two_stage/stages.py | 43 +++++++++++++-- .../two_stage/tests/test_two_stage.py | 52 +++++++++++++++++++ 6 files changed, 143 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9e4b2f9b..419cce6a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Added +- **Two stages: Latent upscaler as an option** (Settings > Engine > Sampling, "How the hand-over is enlarged"). Bicubic stays the default; pick Latent upscaler and a file from models/latent_upscale_models (H3's convolution upscaler) to enlarge the first stage's clean guess with the learned network instead. It runs at the ratio that reaches the full size (it works in even steps, then a bicubic trim fits the exact grid), once per run, on a small latent; the file is loaded once and kept. A missing or unreadable file falls back to bicubic with a note. Untested on a GPU. - **Upscale no longer makes the next generation reload the models.** It ran as an extra ComfyUI prompt without the workflow's loaders, and ComfyUI's cache lets go of loaders a prompt does not use; it now runs as a background job of FunUI, beside ComfyUI's queue (Interrupt in ComfyUI no longer stops it). - **Interpolate** (module of its own, add-on **NVIDIA RTX Frame Generation**): adds in-between frames to a finished render with NVIDIA's RTX frame generation (2 to 4 times the frames; the length and the sound stay, 24 fps becomes 48 fps). Settings > Editor > Interpolate: Never, a button beside the rating, or always when a render finishes. Quality Fast, Balanced or Best, and 10-bit processing. It runs as a background job of FunUI, not as a ComfyUI prompt, so it does not make ComfyUI drop the workflow's loaded models (an extra prompt without loaders does exactly that, and the next generation reloads them); the NVIDIA model also stays loaded between presses. Works beside Upscale in either order and never runs twice on the same file. Needs an RTX card and nvidia-vfx 0.2 or newer: switching the add-on on moves an older nvidia-vfx up (nothing else is touched), and Super Resolution never moves a newer one back. ComfyUI's own RIFE/FILM interpolation nodes are a different thing and are not used. Works on a rental GPU (a few seconds for a clip). - **Project settings > Resolution is the** (H3, LTX image to video): what Width and Height mean once the scene picture sets the shape. Short side (the smaller number is the clip's shorter side), Long side (the larger is its longer side) or Pixel count (the old rule: about Width x Height pixels, which turned 1024x1024 into 1312x768 for a wide picture). New projects start on Short side; projects saved earlier keep Pixel count, so nothing changes size by itself. The H3 preset no longer needs the KJNodes ImageTransformKJ node for this (saved H3 pipelines are switched over on load). diff --git a/README.md b/README.md index 721c3955..29b9e6a7 100644 --- a/README.md +++ b/README.md @@ -168,6 +168,7 @@ I express my deepest gratitude to: - PTD few-step support for MiniMax H3 is ported from [ComfyUI-PTD-Compact](https://github.com/Nynxz/ComfyUI-PTD-Compact) by Nynxz (Apache-2.0), which builds on [ComfyUI-PTD](https://github.com/vita-epfl/PTD-inference) by EPFL VITA Lab (Apache-2.0; paper "Parametric Trajectory Distillation for Few-Step Video Generation", arXiv 2610.11498) and on ComfyUI-MiniMaxH3-PTD by taobao-content-ai (Apache-2.0); licence notice in `modules/sampling/ptd`; - NVIDIA and the NVIDIA Video Effects SDK (`nvidia-vfx`, NVIDIA's own licence) — RTX Super Resolution in Upscale and RTX Frame Generation in Interpolate run on it; FunUI installs it only when you switch the add-on on; - NVIDIA and the authors of Sol-H3 (arXiv 2609.35110, NVlabs/Sana, Apache-2.0) — the MXFP8 loader option (which blocks stay in bf16) and the Two stages setting (the cross-resolution schedule) follow its recipe; +- LBH-123-AI ([Comfyui_Minimax_h3_latent_Upscaler](https://github.com/LBH-123-AI/Comfyui_Minimax_h3_latent_Upscaler)) and TripleHeadedMonkey ([ComfyUI-MiniMaxH3_LatentUpscaler](https://github.com/TripleHeadedMonkey/ComfyUI-MiniMaxH3_LatentUpscaler)) — the H3 latent upscaler network (rebuilt from the published weights, not copied) and its use in Second pass and as the Two stages enlarge option; Without all of you, this project would've been impossible. diff --git a/modules/latent/second_pass/nodes.py b/modules/latent/second_pass/nodes.py index 7a6efe32..ab0e1af7 100644 --- a/modules/latent/second_pass/nodes.py +++ b/modules/latent/second_pass/nodes.py @@ -33,19 +33,24 @@ class FunUILatentUpscalerLoader(io.ComfyNode): @classmethod def execute(cls, upscaler_name: str) -> io.NodeOutput: - path = folder_paths.get_full_path_or_raise(FOLDER, upscaler_name) - sd = comfy.utils.load_torch_file(path, safe_load=True) - for spec, provider in registry_mod.current().providers("latent_upscaler"): - model = provider(sd) - if model is not None: - log.info("FunUI Latent Upscaler", f"{upscaler_name}: {spec.title} upscaler") - return io.NodeOutput(model) - newer = ("It looks like the newer attention-based H3 upscaler, which FunUI cannot build yet: " - "use an older convolution-based H3 upscaler. ") if any(str(k).startswith("base.") for k in sd) else "" - raise RuntimeError( - f"{upscaler_name} is not an upscaler any installed model module can load. {newer}" - f"If it is Lightricks' or Hunyuan's, ComfyUI's own Load Latent Upscale Model " - f"loads it.") + return io.NodeOutput(load_upscaler(upscaler_name)) + + +def load_upscaler(upscaler_name: str): + """The upscaler a file holds, built by whichever model module recognises it.""" + path = folder_paths.get_full_path_or_raise(FOLDER, upscaler_name) + sd = comfy.utils.load_torch_file(path, safe_load=True) + for spec, provider in registry_mod.current().providers("latent_upscaler"): + model = provider(sd) + if model is not None: + log.info("FunUI Latent Upscaler", f"{upscaler_name}: {spec.title} upscaler") + return model + newer = ("It looks like the newer attention-based H3 upscaler, which FunUI cannot build yet: " + "use an older convolution-based H3 upscaler. ") if any(str(k).startswith("base.") for k in sd) else "" + raise RuntimeError( + f"{upscaler_name} is not an upscaler any installed model module can load. {newer}" + f"If it is Lightricks' or Hunyuan's, ComfyUI's own Load Latent Upscale Model " + f"loads it.") def _video_of(samples): diff --git a/modules/sampling/two_stage/__init__.py b/modules/sampling/two_stage/__init__.py index 0a033194..464f2a41 100644 --- a/modules/sampling/two_stage/__init__.py +++ b/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} diff --git a/modules/sampling/two_stage/stages.py b/modules/sampling/two_stage/stages.py index 33cb1e03..9757af2f 100644 --- a/modules/sampling/two_stage/stages.py +++ b/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 diff --git a/modules/sampling/two_stage/tests/test_two_stage.py b/modules/sampling/two_stage/tests/test_two_stage.py index 56e6a940..22ba1373 100644 --- a/modules/sampling/two_stage/tests/test_two_stage.py +++ b/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