From 8d4a199e1fbf404f6acef7f2e81f3d8d3bc72b98 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:38:36 +0300 Subject: [PATCH] feat(sampler): second pass and segmented detailing on any model family MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both were switched off on H3 because Lightricks' upsampler is 128-channel — a fact about the file, not the family. They now load any upsampler ComfyUI can (core's loader handles the rest), refuse a mismatched one by name, and never download the LTX file for H3. --- detailing.py | 73 +++++++++++++++++++++++--- movie_editor/frontend/pipeline_caps.js | 17 +++--- samplers.py | 40 +++++++++----- tests/test_second_pass.py | 72 +++++++++++++++++++++++++ 4 files changed, 171 insertions(+), 31 deletions(-) diff --git a/detailing.py b/detailing.py index 62cab65a..fe15dc86 100644 --- a/detailing.py +++ b/detailing.py @@ -99,7 +99,34 @@ DEFAULT_UPSAMPLER_REPO = "Lightricks/LTX-2.3" DEFAULT_UPSAMPLER_FILE = "ltx-2.3-spatial-upscaler-x2-1.1.safetensors" -def resolve_upsampler_name(model_name=None): +# LTX-2's video latent width. The official Lightricks upsampler is trained on it, so a +# model whose latents are a different width cannot use that file no matter how it is loaded. +LTX_LATENT_CHANNELS = 128 + + +def upsampler_in_channels(model): + """How many latent channels an upsampler expects, or None when it can't be read. + + Used to say "this upsampler is for a different model" instead of letting a shape error + surface from inside a conv three minutes into a render. + """ + for candidate in (model, getattr(model, "model", None)): + if candidate is None: + continue + declared = getattr(candidate, "in_channels", None) + if isinstance(declared, int) and declared > 0: + return declared + state = getattr(candidate, "state_dict", None) + if not callable(state): + continue + for tensor in state().values(): + # The first convolution is the input one; its second dim is what it accepts. + if getattr(tensor, "ndim", 0) >= 4: + return int(tensor.shape[1]) + return None + + +def resolve_upsampler_name(model_name=None, latent_channels=None): """Turn the node's detail_upsampler widget value into a loadable filename. Explicit filename -> returned as-is. "auto" (and the legacy "None" default) @@ -107,6 +134,10 @@ def resolve_upsampler_name(model_name=None): installed latent upscale model, else download the official file from HF. Raises with a user-actionable message when nothing can be obtained — callers surface it, they must NOT silently skip. + + `latent_channels` is the width of the latents this run actually produces. When it is + not LTX's, the official Lightricks file cannot be the answer, so it is never downloaded + on a model it could not serve — a gigabyte for nothing, then a shape error. """ import os @@ -116,6 +147,14 @@ def resolve_upsampler_name(model_name=None): if name and name not in ("None", "auto"): return name files = folder_paths.get_filename_list("latent_upscale_models") + if latent_channels is not None and int(latent_channels) != LTX_LATENT_CHANNELS: + if files: + return files[0] + raise RuntimeError( + f"no latent upscale model installed, and this model's latents are " + f"{int(latent_channels)}-channel, so the official Lightricks upsampler " + f"({LTX_LATENT_CHANNELS}-channel) would not fit it — put an upsampler trained " + "for this model in models/latent_upscale_models") if DEFAULT_UPSAMPLER_FILE in files: return DEFAULT_UPSAMPLER_FILE spatial = [f for f in files if "spatial" in f.lower() and "upscal" in f.lower()] @@ -166,8 +205,9 @@ def load_latent_upsampler(model_name): model_path = folder_paths.get_full_path_or_raise("latent_upscale_models", model_name) sd, metadata = comfy.utils.load_torch_file(model_path, safe_load=True, return_metadata=True) if "post_upsample_res_blocks.0.conv2.bias" not in sd: - raise ValueError( - f"{model_name} is not an LTX latent upsampler (expected post_upsample_res_blocks keys).") + # Not Lightricks' architecture. ComfyUI's own loader knows the others (and will know + # the next one), so hand it over rather than declaring the file unusable here. + return _load_upsampler_via_core(model_name) config = json.loads(metadata["config"]) # ComfyUI v0.29.0 (upstream f8a3fd9d) moved the upsampler onto DynamicVram and made # `operations` a required argument of from_config; older cores build nn.Modules @@ -185,6 +225,23 @@ def load_latent_upsampler(model_name): return model +def _load_upsampler_via_core(model_name): + """ComfyUI's Load Latent Upscale Model, for every architecture that is not LTX's.""" + try: + from comfy_extras.nodes_hunyuan import LatentUpscaleModelLoader + except ImportError as exc: + raise ValueError( + f"{model_name} is not an LTX latent upsampler, and this ComfyUI has no other " + "latent upscale model loader to try") from exc + out = LatentUpscaleModelLoader.execute(model_name) + result = getattr(out, "result", out) + if isinstance(result, (tuple, list)): + result = result[0] + if result is None: + raise ValueError(f"{model_name} could not be loaded as a latent upscale model") + return result + + def _get_clipseg(): """CLIPSeg processor+model, loaded once and kept on CPU. @@ -380,21 +437,23 @@ def _run_upsampler(upsampler, video_crop, vae, debug=False): """ import comfy.model_management + # Core's loader hands back a patcher; FunPack's LTX path hands back the module itself. + module = getattr(upsampler, "model", upsampler) device = comfy.model_management.get_torch_device() - model_dtype = next(upsampler.parameters()).dtype + model_dtype = next(module.parameters()).dtype in_dtype, in_device = video_crop.dtype, video_crop.device stats = _per_channel_stats(vae) try: - upsampler.to(device) + module.to(device) x = video_crop.to(device=device, dtype=model_dtype) if stats is not None: x = stats.un_normalize(x) with torch.no_grad(): - x = upsampler(x) + x = module(x) if stats is not None: x = stats.normalize(x) finally: - upsampler.cpu() + module.cpu() return x.to(device=in_device, dtype=in_dtype) diff --git a/movie_editor/frontend/pipeline_caps.js b/movie_editor/frontend/pipeline_caps.js index adf2e360..3ce6d68f 100644 --- a/movie_editor/frontend/pipeline_caps.js +++ b/movie_editor/frontend/pipeline_caps.js @@ -128,20 +128,15 @@ bounded_attention_enabled: "Bounded Attention masks text cross-attention; H3 puts text in the same self-attention " + "stream as the video, where the mask would be far too large to hold.", - segmented_detailing: - "Segmented detailing uses Lightricks' latent upsampler, trained on LTX's 128-channel " - + "latent; H3's latent has 24 channels.", }; // Same idea for non-boolean knobs whose neutral value means "off". - function h3DeadValueIssues(si) { - const out = []; - if (si.second_pass && si.second_pass_op && si.second_pass_op !== "none") { - out.push(["second_pass_op", - "The second pass still runs on H3, but its latent op ('" + si.second_pass_op + "') uses the " - + "same LTX-only upsampler, so the op is skipped."]); - } - return out; + // The latent ops (second_pass_op, segmented detailing) are NOT listed here any more: what + // they need is an upsampler whose latents are the same width as the model's, which depends + // on what is installed rather than on the family, and the sampler reports a mismatch by + // name. A chip here would be a guess about the user's models folder. + function h3DeadValueIssues(_si) { + return []; } // The mirror image of H3_DEAD_SAMPLER_INPUTS: settings that only mean something ON H3. diff --git a/samplers.py b/samplers.py index 6d8edadf..0a2fc272 100644 --- a/samplers.py +++ b/samplers.py @@ -2788,7 +2788,7 @@ class FunPackLTXAVSceneChainSampler: }), "segmented_detailing": ("BOOLEAN", { "default": False, - "tooltip": "EXPERIMENTAL ADetailer-for-video: after each scene finishes denoising, CLIPSeg (text-prompted segmentation) locates the regions named in detail_targets on a few decoded keyframes; the matched region is cut out of the video latent as a spatiotemporal tube, pushed through Lightricks' trained latent upsampler (2x spatial — the official two-stage pipeline's stage-2 model), then either used directly (detail_mode='sharpen', near-free) or re-noised + re-denoised for a 3-step tail (detail_mode='repair', default — costs ~4x tube area fraction x 3 steps, hands ~+15%). Downscaled back to its ORIGINAL latent size and pasted through the feathered CLIPSeg silhouette either way. Final resolution never changes. Tubes over detail_max_area (default 35%) are refused as a cost guard, not a content judgment. Audio untouched by construction. detail_upsampler 'auto' finds or downloads the official Lightricks upsampler (~1 GB, once); skips are reported loudly in console + scene report. UNVALIDATED LIVE.", + "tooltip": "EXPERIMENTAL ADetailer-for-video: after each scene finishes denoising, CLIPSeg (text-prompted segmentation) locates the regions named in detail_targets on a few decoded keyframes; the matched region is cut out of the video latent as a spatiotemporal tube, pushed through Lightricks' trained latent upsampler (2x spatial — the official two-stage pipeline's stage-2 model), then either used directly (detail_mode='sharpen', near-free) or re-noised + re-denoised for a 3-step tail (detail_mode='repair', default — costs ~4x tube area fraction x 3 steps, hands ~+15%). Downscaled back to its ORIGINAL latent size and pasted through the feathered CLIPSeg silhouette either way. Final resolution never changes. Tubes over detail_max_area (default 35%) are refused as a cost guard, not a content judgment. Audio untouched by construction. detail_upsampler 'auto' finds or downloads the official Lightricks upsampler (~1 GB, once) when the model's latents are LTX-width, and otherwise uses an installed upsampler without downloading anything; skips are reported loudly in console + scene report. UNVALIDATED LIVE.", }), "detail_targets": ("STRING", { "default": "hands", @@ -2863,7 +2863,7 @@ class FunPackLTXAVSceneChainSampler: }), "second_pass_op": (["none", "sharpen", "upscale_2x"], { "default": "none", - "tooltip": "OPTIONAL latent-space operation applied between the two passes — 'none' by default, nothing runs unless you pick one. 'sharpen': one forward of Lightricks' trained 2x latent upsampler, resampled straight back to the original size. No video-model calls at all, so it costs a fraction of a step; pass 2 then re-denoises the sharpened latent, which is what makes it stick. It adds detail consistent with what is already there and CANNOT fix structure that is wrong (an extra finger stays an extra finger, just sharper) — the same limit segmented detailing's sharpen mode documents. 'upscale_2x': the same upsampler, but the result is KEPT at 2x, so pass 2 runs at four times the pixels and the scene decodes at double resolution. That is 3-5x the sampling cost of the second half. The i2v pin SURVIVES it: the pinned frames are carried through the upsampler with everything else and the mask is scaled to the new grid, so pass 2 still holds the anchor — as the upscaled anchor rather than the encoded source image, which the scene report states. Guide keyframes do not survive, because they are token indices into the old grid; pass 1 uses them in full. MULTI-SCENE works with it: a scene finishes at 2x while every later scene is still built from the latent template at the original size, so everything that crosses a scene boundary (carried overlap frames, the anchor's continuation, the soft join, JoyAI memory, per-scene guide sources) is brought back to the template's grid on the way. Each scene still samples and OUTPUTS at 2x — only the carried material is resampled, and only downwards, which is the direction that survives it: those frames exist to say 'continue from here', which a resample preserves far better than invented detail would. Both ops use the same upsampler file as segmented detailing (detail_upsampler, 'auto' downloads the official one on first use). Video stream only; audio is never reshaped.", + "tooltip": "OPTIONAL latent-space operation applied between the two passes — 'none' by default, nothing runs unless you pick one. 'sharpen': one forward of Lightricks' trained 2x latent upsampler, resampled straight back to the original size. No video-model calls at all, so it costs a fraction of a step; pass 2 then re-denoises the sharpened latent, which is what makes it stick. It adds detail consistent with what is already there and CANNOT fix structure that is wrong (an extra finger stays an extra finger, just sharper) — the same limit segmented detailing's sharpen mode documents. 'upscale_2x': the same upsampler, but the result is KEPT at 2x, so pass 2 runs at four times the pixels and the scene decodes at double resolution. That is 3-5x the sampling cost of the second half. The i2v pin SURVIVES it: the pinned frames are carried through the upsampler with everything else and the mask is scaled to the new grid, so pass 2 still holds the anchor — as the upscaled anchor rather than the encoded source image, which the scene report states. Guide keyframes do not survive, because they are token indices into the old grid; pass 1 uses them in full. MULTI-SCENE works with it: a scene finishes at 2x while every later scene is still built from the latent template at the original size, so everything that crosses a scene boundary (carried overlap frames, the anchor's continuation, the soft join, JoyAI memory, per-scene guide sources) is brought back to the template's grid on the way. Each scene still samples and OUTPUTS at 2x — only the carried material is resampled, and only downwards, which is the direction that survives it: those frames exist to say 'continue from here', which a resample preserves far better than invented detail would. Both ops use the same upsampler file as segmented detailing (detail_upsampler, 'auto' downloads the official Lightricks one on first use). Works on any model family, not just LTX — what it needs is an upsampler whose latents are the same width as this model's, so on MiniMax H3 install an H3 latent upsampler and pick it here; 'auto' will not download the LTX file for a model it cannot fit, and a mismatch is reported as a skip with both channel counts rather than failing the render. Video stream only; audio is never reshaped.", }), "h3_audio_clock": ("BOOLEAN", { "default": False, @@ -3758,6 +3758,14 @@ class FunPackLTXAVSceneChainSampler: tensors = self._latent_tensors(result) video = tensors[0] h, w = int(video.shape[-2]), int(video.shape[-1]) + # An upsampler trained on a different model's latents fails inside a convolution, + # minutes in, with a shape error that names neither model. Say it here instead. + want = _d.upsampler_in_channels(upsampler) + have = int(video.shape[1]) if video.ndim >= 5 else None + if want and have and want != have: + return latent, (f"second_pass_op={op} skipped: the selected latent upsampler takes " + f"{want}-channel latents and this model's are {have}-channel — " + f"install an upsampler trained for this model") try: up = _d._run_upsampler(upsampler, video, vae) except Exception as exc: # noqa: BLE001 @@ -6272,12 +6280,6 @@ class FunPackLTXAVSceneChainSampler: and abs(float(v2a_grad_scale) - 1.0) > 1e-6): _dead.append("v2a_grad_scale (hooks LTXAV's video_to_audio_attn submodule; H3 " "has no separate video->audio cross-attention to scale)") - if segmented_detailing: - _dead.append("segmented_detailing (Lightricks' latent upsampler is trained on " - "LTX's 128-channel latent, not H3's 24-channel one)") - if second_pass and str(second_pass_op or "none").lower() in ("sharpen", "upscale_2x"): - _dead.append(f"second_pass_op='{second_pass_op}' (same LTX-only latent upsampler; " - "the second pass itself still runs, just without the latent op)") for _line in _dead: print(f"[FunPackSceneChain] H3: {_line} — SKIPPED.") # Turn them off for real rather than letting each one discover its own missing @@ -6285,10 +6287,10 @@ class FunPackLTXAVSceneChainSampler: # that dies three minutes in is worse than a knob that says why it is inert. bounded_attention_enabled = False identity_transfer_enabled = False - segmented_detailing = False v2a_grad_scale = 1.0 - if second_pass and str(second_pass_op or "none").lower() != "none": - second_pass_op = "none" + # second_pass_op is NOT switched off here. It only needs a latent upsampler that + # matches this model's latent width; which file that is depends on what is + # installed, not on the family, so it is checked where the upsampler is loaded. # Defensively strip any enhancement block hooks left on the shared diffusion # model by a previous run (build_enhancements only removes them on scene @@ -6949,7 +6951,13 @@ class FunPackLTXAVSceneChainSampler: from . import detailing as _det except ImportError: import detailing as _det - _r = _det.resolve_upsampler_name(detail_upsampler) + _chans = None + try: + _v = self._latent_tensors(_sp_state)[0] + _chans = int(_v.shape[1]) if _v.ndim >= 5 else None + except Exception: # noqa: BLE001 + pass + _r = _det.resolve_upsampler_name(detail_upsampler, _chans) _detail_upsampler_model = _det.load_latent_upsampler(_r) except Exception as _exc: # noqa: BLE001 _detail_disabled_reason = str(_exc) @@ -7093,7 +7101,13 @@ class FunPackLTXAVSceneChainSampler: except ImportError: import detailing as _detailing if _detail_upsampler_model is None: - _resolved = _detailing.resolve_upsampler_name(detail_upsampler) + _chans = None + try: + _v = self._latent_tensors(sampled)[0] + _chans = int(_v.shape[1]) if _v.ndim >= 5 else None + except Exception: # noqa: BLE001 + pass + _resolved = _detailing.resolve_upsampler_name(detail_upsampler, _chans) _detail_upsampler_model = _detailing.load_latent_upsampler(_resolved) print(f"[FunPackSceneChain] latent upsampler loaded: {_resolved}") _t_detail0 = _time.perf_counter() diff --git a/tests/test_second_pass.py b/tests/test_second_pass.py index 815829a3..6e3e9b29 100644 --- a/tests/test_second_pass.py +++ b/tests/test_second_pass.py @@ -219,3 +219,75 @@ def test_the_pin_is_dropped_when_the_mask_cannot_be_rebuilt(): chunk = {"samples": clean, "noise_mask": mask} upscaled = {"samples": torch.zeros(1, 4, 4, 16, 16)} assert _node()._restore_pinned_prefix(upscaled, chunk) is upscaled + + +# ── a latent op is about the upsampler, not the model family ────────────────── +# The op used to be switched off outright on MiniMax H3 because Lightricks' upsampler is +# trained on LTX's 128-channel latent. That is a fact about the FILE, not about H3: with an +# upsampler trained for the model, the op runs. What has to stay true is that a mismatched +# one is refused by name rather than blowing up inside a convolution mid-render. +import detailing # noqa: E402 + + +class _TypedUpsampler(_FakeUpsampler): + def __init__(self, in_channels): + self.in_channels = in_channels + + +def test_an_upsamplers_width_is_read_from_the_model(): + assert detailing.upsampler_in_channels(_TypedUpsampler(24)) == 24 + + +def test_an_upsamplers_width_falls_back_to_its_first_convolution(): + class _Weights: + def state_dict(self): + return {"initial_conv.weight": torch.zeros(512, 24, 3, 3)} + assert detailing.upsampler_in_channels(_Weights()) == 24 + + +def test_an_upsampler_wrapped_in_a_patcher_is_still_readable(): + class _Patcher: + model = _TypedUpsampler(24) + assert detailing.upsampler_in_channels(_Patcher()) == 24 + + +def test_an_upsampler_for_another_model_is_refused_by_name(): + lat = _video() # 4-channel stand-in latents + out, note = _node()._second_pass_operate(lat, "upscale_2x", _TypedUpsampler(128), None) + assert out is lat + assert "128-channel" in note and "4-channel" in note + + +def test_a_matching_upsampler_runs_on_any_family(): + out, note = _node()._second_pass_operate(_video(h=8, w=8), "upscale_2x", _TypedUpsampler(4), None) + assert out["samples"].shape == (1, 4, 4, 16, 16) + assert "4x the pixels" in note + + +def _installed(monkeypatch, files): + """Patch the SHARED folder_paths stub in place. Replacing sys.modules["folder_paths"] + would leak this test's view of the models folder into every later test file.""" + monkeypatch.setattr(sys.modules["folder_paths"], "get_filename_list", + lambda _kind: list(files)) + + +def test_auto_never_downloads_the_ltx_upsampler_for_a_model_it_cannot_fit(monkeypatch): + """A gigabyte, then a shape error. 'auto' takes what is installed instead.""" + _installed(monkeypatch, ["h3_latent_upscaler.safetensors"]) + monkeypatch.setattr(detailing, "DEFAULT_UPSAMPLER_FILE", "ltx.safetensors", raising=False) + assert detailing.resolve_upsampler_name("auto", 24) == "h3_latent_upscaler.safetensors" + + +def test_auto_says_what_to_install_when_nothing_fits(monkeypatch): + _installed(monkeypatch, []) + try: + detailing.resolve_upsampler_name("auto", 24) + except RuntimeError as exc: + assert "24-channel" in str(exc) and "latent_upscale_models" in str(exc) + else: + raise AssertionError("expected a RuntimeError naming the channel width") + + +def test_an_explicitly_chosen_file_is_always_honoured(monkeypatch): + _installed(monkeypatch, []) + assert detailing.resolve_upsampler_name("my_upsampler.safetensors", 24) == "my_upsampler.safetensors"