From 0ed59f06d2f77509ac15f37a35b2761942e8842e Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Thu, 5 Nov 2026 20:19:04 +0200 Subject: [PATCH] Two stages: Trilinear removed (migrated to Bicubic via renamed_options); Save Video: opt-in edge smoothing Co-Authored-By: Claude Sonnet 5.5 --- CHANGELOG.md | 3 + README.md | 1 + core/schema.py | 6 ++ core/tests/test_schema.py | 23 ++++++++ core/tests/test_shipped_nodes.py | 2 +- modules/output/save/nodes.py | 36 ++++++++++- modules/output/save/tests/test_save.py | 59 +++++++++++++++++++ modules/sampling/two_stage/__init__.py | 5 +- modules/sampling/two_stage/stages.py | 5 +- .../two_stage/tests/test_two_stage.py | 2 +- 10 files changed, 131 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 09ada0ba..1422b515 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,8 @@ - **Video input is the project's only.** The per-scene copy and the "Reference videos use" choice are gone; scenes' old own settings are no longer read. - **Continuity moved into Project settings** as "Pass the last frame of previous scene as reference into next scene". The Continuity page left Settings; "Don't continue from a fade to black" is no longer a switch, a shot after a clip that ends dark always starts without a picture. +- **Two stages: "Trilinear" is gone from the enlarge methods.** It gave exactly the same result as Bilinear (time is never resampled), so it was a duplicate with a misleading name. A project or setting that saved it now loads as Bicubic, the default (any enum setting can declare `renamed_options` for this). + ### Fixed - **No more torchsde "Should have ta>=t0" warning** on SDE-type samplers (it appeared whenever the schedule's smallest sigma had a 7th decimal of 5 or more, e.g. with Two stages). ComfyUI's Brownian noise now starts a hair earlier and ends a hair later than the schedule, so torchsde's rounding never lands inside it. The noise itself is unchanged and the picture was never affected. @@ -18,6 +20,7 @@ ### Added - **Run meter** (Settings > Engine > Sampling, off by default). Logs how grainy the picture is, across it and over time, and how much it still changes at each step, then compares with the previous run of the same seed. It adds no model call, so flipping one setting and rendering again shows what that setting did. +- **Save Video: "smoothing"** (off by default; light or strong). Anti-aliasing for jagged edges in the FXAA idea: only pixels on a strongly contrasting edge are blended with their neighbours, flat areas and texture stay, frames never mix. It runs on the frames' own device before encoding, so NVENC, its CPU fallback and every codec get the same picture. Tried on the CPU only. - **LoRA Loader: "Apply to stage"** (First, Second or Both, per slot). Shown only while the Two stages module is present and not switched off for the project; the feature itself may stay off in the sampler settings. A LoRA set to First works only in the smaller first stage, one set to Second only in the full-size steps; Two stages hands each stage the model with its own LoRAs, and ComfyUI re-patches the weights between the stages. Only the model changes for First or Second (the text encoder is encoded before sampling, so its strength is not used). When Two stages is off, or a run cannot be cut in two, those LoRAs are not applied and the Log says so. Both is the old behaviour. Tried on a small model on the CPU with every feature switched on; not on a GPU yet. - **LoRA Loader reads LoRAs written for the Diffusers layout of H3** (separate q, k and v, `ff.net`, `transformer_blocks`), such as NVIDIA's Open-H3-Regenerate-2K refiner, which used to report "matched nothing". The names are translated, q, k and v are aimed at their rows of H3's fused attention weight, and the halves of the gated MLP are put in H3's order (its layout stores them the other way round, so without that the LoRA loads and quietly does the wrong thing). Strength 1.0, as the author says it is already folded in. - **Attention explorer: "Change the sound too"** (Settings > Engine > Sampling). Turn it off and only the picture's attention changes: the sound's rows keep the model's own attention and the sound is no longer one of the clip regions. On models that mix picture and sound in one sequence, Prompt focus stays off while the sound is kept as is. On by default, so nothing changes until you switch it. Untested on a GPU. diff --git a/README.md b/README.md index 88dd092a..a6eb6fca 100644 --- a/README.md +++ b/README.md @@ -171,6 +171,7 @@ I express my deepest gratitude to: - Tero Karras, Miika Aittala, Timo Aila and Samuli Laine (Elucidating the Design Space of Diffusion-Based Generative Models, arXiv 2206.00364) — the churn idea (fresh noise added during sampling) behind the Turbo sampler's fresh-noise step; - Don Mitchell and Arun Netravali (the cubic filter family, SIGGRAPH 1988), Edwin Catmull and Raphael Rom (the Catmull-Rom spline, 1974) and Claude Duchon (the Lanczos window, 1979) — the Mitchell, Catmull-Rom and Lanczos choices when Two stages resizes the latent; - NVIDIA and the authors of Open-H3-Regenerate-2K (NVlabs/Sana) — the refiner LoRA's published layout, sampling recipe (a second stage at the full size) and strength convention that the LoRA Loader's Diffusers-layout reading and its "Apply to stage" option are built for; and Kijai, whose published conversion notes (fused attention, order of the gated MLP halves) were used to check our translation; +- Timothy Lottes (NVIDIA, FXAA) - the idea behind Save Video smoothing: soften only the pixels on a contrasty edge (written fresh, not copied); - 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; - The Turbo sampler (stochastic flow step, polish step) rests on consistency models (Song, Dhariwal, Chen, Sutskever; arXiv 2303.01469), latent consistency models (Luo, Tan, Huang, Li, Zhao; arXiv 2310.04378) and the ancestral noise injection of k-diffusion by Katherine Crowson; its second-order late steps reuse Distilled Flow; diff --git a/core/schema.py b/core/schema.py index 67582de6..975f5cbe 100644 --- a/core/schema.py +++ b/core/schema.py @@ -92,6 +92,10 @@ def _validate_enum(spec: dict, where: str) -> None: raise SchemaError(f"{where} has duplicate option values.") if spec["default"] not in values: raise SchemaError(f"{where}'s default {spec['default']!r} is not one of its options.") + # {old value: current option}: an option that was removed, mapped forward when saved settings are read. + for old, new in (spec.get("renamed_options") or {}).items(): + if old in values or new not in values: + raise SchemaError(f"{where}'s renamed_options {old!r} -> {new!r} must map a removed value to a current option.") # Accepted as #rgb, #rrggbb or #rrggbbaa. One spelling in, one spelling out -- @@ -272,6 +276,8 @@ def check_values(spec, values: Any, keep_bad: bool = False) -> Tuple[Dict[str, A log.once(f"stale_setting:{spec.id}.{key}", log.WARNING, "FunUI settings", f"{spec.id}: saved setting {key!r} ignored: this module has no such setting") continue + if isinstance(value, str): + value = (declared.get("renamed_options") or {}).get(value, value) problem = _value_problem(declared, key, value, spec.id) if problem and not _applies(declared, {**spec.defaults(), **values}): continue # a setting that is hidden by another one's choice cannot act, so it cannot stop the run: its default stands diff --git a/core/tests/test_schema.py b/core/tests/test_schema.py index f3867f19..f38a1ce6 100644 --- a/core/tests/test_schema.py +++ b/core/tests/test_schema.py @@ -343,3 +343,26 @@ def test_a_hidden_setting_with_a_bad_value_does_not_stop_the_run(): assert problems == [] and clean["file"] == "a.pth" clean, problems = schema.check_values(spec, {"method": "file", "file": "other.safetensors"}) assert problems and "m.file must be one of" in problems[0] + + +def test_a_removed_enum_option_is_mapped_forward_when_saved_settings_are_read(): + from types import SimpleNamespace + from core import schema + s = {"m": {"type": "enum", "default": "a", "options": [{"value": "a", "label": "A"}, {"value": "b", "label": "B"}], "renamed_options": {"gone": "b"}}} + spec = SimpleNamespace(id="x", settings=s, defaults=lambda: {"m": "a"}) + assert schema.check_values(spec, {"m": "gone"}) == ({"m": "b"}, []) + assert schema.check_values(spec, {"m": "other"})[1] # an unknown value is still a problem + with pytest.raises(schema.SchemaError): + schema.validate_setting("m", {**s["m"], "renamed_options": {"a": "b"}}, s) # a live option cannot be "removed" + with pytest.raises(schema.SchemaError): + schema.validate_setting("m", {**s["m"], "renamed_options": {"gone": "zzz"}}, s) + + +def test_two_stages_saved_trilinear_loads_as_the_default_method(): + from core import schema + from modules.sampling import two_stage + spec = two_stage.SETTINGS["resize"] + assert "trilinear" not in [o["value"] for o in spec["options"]] + fake = type("S", (), {"id": "two_stage", "settings": two_stage.SETTINGS, "defaults": lambda self: {k: v["default"] for k, v in two_stage.SETTINGS.items()}})() + clean, problems = schema.check_values(fake, {"enabled": True, "resize": "trilinear"}) + assert problems == [] and clean["resize"] == "bicubic" diff --git a/core/tests/test_shipped_nodes.py b/core/tests/test_shipped_nodes.py index 04a8a915..972df472 100644 --- a/core/tests/test_shipped_nodes.py +++ b/core/tests/test_shipped_nodes.py @@ -54,7 +54,7 @@ FROZEN = { "FunUILoadModifiers": ["model", "settings"], "FunUILatentUpscalerLoader": ["upscaler_name"], "FunUICutOpening": ["images", "frames", "fps", "audio"], - "FunUISaveVideo": ["images", "audio", "fps", "filename_prefix", "encoder", "codec", "quality"], + "FunUISaveVideo": ["images", "audio", "fps", "filename_prefix", "encoder", "codec", "quality", "smoothing"], "FunUIPromptMarkup": ["text"], "FunUINegativeErase": ["positive", "negative", "strength", "mode", "keep_size"], "FunUIShotCamera": ["text", "settings", "length", "frame_rate", "project"], diff --git a/modules/output/save/nodes.py b/modules/output/save/nodes.py index 5cc90077..4a566955 100644 --- a/modules/output/save/nodes.py +++ b/modules/output/save/nodes.py @@ -15,6 +15,9 @@ from fractions import Fraction import av import folder_paths +import numpy as np +import torch +import torch.nn.functional as F from comfy.cli_args import args from comfy_api.latest import io, ui @@ -31,6 +34,31 @@ QUALITY = {"nvenc": {"standard": "23", "high": "16", "maximum": "10"}, NVENC = {"h265": ("hevc_nvenc", {"preset": "p6", "rc": "vbr", "cq": "16", "b:v": "0"}), "h264": ("h264_nvenc", {"preset": "p6", "rc": "vbr", "cq": "16", "b:v": "0"})} X264 = ("libx264", {"crf": "16", "preset": "medium"}) +# Edge smoothing (the FXAA idea): only pixels on a contrasty edge are blended with their 3x3 neighbours, flat areas and +# texture stay as they are. How much of the blend is kept is `smoothing`. Done on the frames' own device before the +# encoder, so NVENC, its CPU fallback and every codec see the same picture; spatial only, so frames never smear into each other. +SMOOTHINGS = ("off", "light", "strong") +_AMOUNT = {"light": 0.5, "strong": 1.0} +_CHUNK = 8 # frames per step: the float working copy stays small on a long 4K clip + + +def smooth(x, level): + """x: [T, H, W, 3] floats in 0..1 -> the same shape, dtype and device with edges anti-aliased (`level` "off": x itself).""" + amount = _AMOUNT.get(level) + if amount is None: + return x + f = x.permute(0, 3, 1, 2).float() + luma = (f[:, 0:1] * 0.299 + f[:, 1:2] * 0.587 + f[:, 2:3] * 0.114) + p = F.pad(luma, (1, 1, 1, 1), mode="replicate") + contrast = F.max_pool2d(p, 3, 1) + F.max_pool2d(-p, 3, 1) # local max - local min + edge = ((contrast - 0.04) / 0.16).clamp(0, 1) * amount + blur = F.avg_pool2d(F.pad(f, (1, 1, 1, 1), mode="replicate"), 3, 1) + return (f + (blur - f) * edge).permute(0, 2, 3, 1).to(x.dtype) + + +def to_frames(images, smoothing="off"): + """-> uint8 [T, H, W, 3] on the CPU. Rounded: a bare cast cuts every value down, a steady half-level darkening.""" + return np.concatenate([(smooth(x, smoothing) * 255).round().clamp(0, 255).byte().cpu().numpy() for x in images.split(_CHUNK)]) def _with_quality(codec, quality): @@ -138,6 +166,8 @@ class FunUISaveVideo(io.ComfyNode): tooltip="For the GPU encoder. H.264: plays in every browser. H.265: smaller files, but Firefox may not play it."), io.Combo.Input("quality", options=list(QUALITIES), default="high", optional=True, tooltip="How closely the file keeps the picture. standard: small files, but fine texture smears and blocks show on a sharp high-resolution picture. high: visually clean for a bigger file. maximum: near lossless, large."), + io.Combo.Input("smoothing", options=list(SMOOTHINGS), default="off", optional=True, + tooltip="Anti-aliasing for jagged edges (like FXAA in games): only edges with strong contrast are softened, flat areas and texture stay. off: untouched. light: half strength. strong: full strength. A small cost on the GPU, a noticeable one on the CPU for 4K clips."), ], hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo], is_output_node=True, @@ -145,8 +175,8 @@ class FunUISaveVideo(io.ComfyNode): ) @classmethod - def execute(cls, images, fps=24.0, filename_prefix="FunUI", audio=None, encoder="auto", codec="h264", quality="high") -> io.NodeOutput: - frames = (images * 255).round().clamp(0, 255).byte().cpu().numpy() # rounded: a bare cast cuts every value down, a steady half-level darkening + def execute(cls, images, fps=24.0, filename_prefix="FunUI", audio=None, encoder="auto", codec="h264", quality="high", smoothing="off") -> io.NodeOutput: + frames = to_frames(images, smoothing) width, height = frames.shape[2], frames.shape[1] folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path( filename_prefix, folder_paths.get_output_directory(), width, height) @@ -156,6 +186,6 @@ class FunUISaveVideo(io.ComfyNode): file = f"{filename}_{counter:05}_.mp4" t = time.perf_counter() used = encode(os.path.join(folder, file), frames, float(fps), audio, metadata, encoder, codec, quality) - log.info("FunUI Save Video", f"{file}: {len(frames)} frames in {time.perf_counter() - t:.1f}s with {used}, quality {quality}") + log.info("FunUI Save Video", f"{file}: {len(frames)} frames in {time.perf_counter() - t:.1f}s with {used}, quality {quality}" + (f", smoothing {smoothing}" if smoothing != "off" else "")) keep_last_frame(file, frames[-1]) return io.NodeOutput(ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)])) diff --git a/modules/output/save/tests/test_save.py b/modules/output/save/tests/test_save.py index 3e67b331..75948570 100644 --- a/modules/output/save/tests/test_save.py +++ b/modules/output/save/tests/test_save.py @@ -109,3 +109,62 @@ def test_a_failed_keep_never_leaves_an_older_videos_picture_under_the_name(tmp_p monkeypatch.setattr("PIL.Image.Image.save", lambda *a, **k: (_ for _ in ()).throw(OSError("disk full"))) nodes.keep_last_frame("FunUI_00001_.mp4", np.zeros((4, 4, 3), dtype=np.uint8)) assert not old.exists() + + +# --- smoothing (edge anti-aliasing) ---------------------------------------------------------------------------------- + +def _staircase(): + """A hard diagonal edge (black / white), the picture a jagged edge makes: [T, H, W, 3] floats.""" + import torch + yy, xx = torch.meshgrid(torch.arange(32), torch.arange(32), indexing="ij") + return (xx > yy).float()[None, :, :, None].repeat(2, 1, 1, 3) + + +def test_smoothing_off_is_the_old_path_and_the_edge_is_softened_only_where_it_is(): + import torch + x = _staircase() + assert nodes.smooth(x, "off") is x + light, strong = nodes.smooth(x, "light"), nodes.smooth(x, "strong") + assert ((light - x).abs() > 0.05).any() and ((strong - x).abs() > (light - x).abs()).any() + assert torch.equal(light[:, 0, -1], x[:, 0, -1]) and torch.equal(light[:, -1, 0], x[:, -1, 0]) # flat corners untouched + flat = torch.full((1, 16, 16, 3), 0.4) + assert torch.allclose(nodes.smooth(flat, "strong"), flat) + assert torch.equal(light[0], light[1]) # frames never mix: spatial only + + +def test_smoothing_follows_the_device_and_dtype_of_its_input(): + import torch + for dtype in (torch.float16, torch.bfloat16, torch.float32): + x = _staircase().to(dtype) + assert nodes.smooth(x, "strong").dtype == dtype + meta = torch.empty(2, 32, 32, 3, device="meta", dtype=torch.float16) # the cloud stand-in for "some other device" + out = nodes.smooth(meta, "light") + assert out.device.type == "meta" and out.dtype == torch.float16 and out.shape == meta.shape + + +def test_chunked_conversion_equals_one_go_and_off_equals_the_plain_cast(): + import torch + x = torch.rand(nodes._CHUNK * 2 + 3, 16, 16, 3) + plain = (x * 255).round().clamp(0, 255).byte().numpy() + assert (nodes.to_frames(x, "off") == plain).all() + whole = (nodes.smooth(x, "strong") * 255).round().clamp(0, 255).byte().numpy() + assert (nodes.to_frames(x, "strong") == whole).all() + + +def test_every_smoothing_saves_with_every_encoder_codec_and_quality(tmp_path): + # registry-driven: the cross product of what the node offers; no NVENC here, so "auto"/"gpu" are the NVENC -> CPU fallback + import itertools + x = _staircase() + for smoothing, encoder, codec, quality in itertools.product(nodes.SMOOTHINGS, nodes.ENCODERS, nodes.CODECS, nodes.QUALITIES): + path = str(tmp_path / f"{smoothing}_{encoder}_{codec}_{quality}.mp4") + frames = nodes.to_frames(x, smoothing) + used = nodes.encode(path, frames, 24.0, encoder=encoder, codec=codec, quality=quality) + assert used and os.path.getsize(path) > 0 + with av.open(path) as f: + assert f.streams.video[0].width == 32 + + +def test_smoothing_is_a_node_input_defaulting_to_off(): + schema = nodes.FunUISaveVideo.define_schema() + combo = next(i for i in schema.inputs if i.id == "smoothing") + assert combo.default == "off" and list(combo.options) == list(nodes.SMOOTHINGS) diff --git a/modules/sampling/two_stage/__init__.py b/modules/sampling/two_stage/__init__.py index 61074f21..1bd5ea15 100644 --- a/modules/sampling/two_stage/__init__.py +++ b/modules/sampling/two_stage/__init__.py @@ -69,11 +69,12 @@ SETTINGS = { "options": [{"value": "bicubic", "label": "Bicubic"}, {"value": "mitchell", "label": "Mitchell"}, {"value": "catmull-rom", "label": "Catmull-Rom"}, {"value": "lanczos", "label": "Lanczos"}, {"value": "bislerp", "label": "Bislerp"}, {"value": "bilinear", "label": "Bilinear"}, - {"value": "trilinear", "label": "Trilinear"}, {"value": "area", "label": "Area"}, + {"value": "area", "label": "Area"}, {"value": "nearest-exact", "label": "Nearest-exact"}, {"value": "upscaler", "label": "Latent upscaler"}], + "renamed_options": {"trilinear": "bicubic"}, # removed: it was Bilinear here (time is never resampled) "label": "How the hand-over is enlarged", - "hint": "The plain methods are free (Trilinear gives the same as Bilinear here: time is never resampled). Latent upscaler runs H3's learned latent upscaler (models/latent_upscale_models) " + "hint": "The plain methods are 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}, diff --git a/modules/sampling/two_stage/stages.py b/modules/sampling/two_stage/stages.py index 601d3ccd..29c5969c 100644 --- a/modules/sampling/two_stage/stages.py +++ b/modules/sampling/two_stage/stages.py @@ -74,7 +74,7 @@ def _upscaler(name): return _CACHE["model"] -METHODS = ("bicubic", "mitchell", "catmull-rom", "lanczos", "bislerp", "bilinear", "trilinear", "area", "nearest-exact") +METHODS = ("bicubic", "mitchell", "catmull-rom", "lanczos", "bislerp", "bilinear", "area", "nearest-exact") def _bc_kernel(B, C): @@ -123,9 +123,6 @@ def resize_video(video, h, w, method="bicubic"): return downscale_to(video, h, w) if method in KERNELS: return _filtered(video.float(), h, w, method).to(video.dtype) - if method == "trilinear": - return F.interpolate(video.float(), size=(video.shape[2], h, w), mode="trilinear", - align_corners=False).to(video.dtype) b, c, f, hh, ww = video.shape x = video.permute(0, 2, 1, 3, 4).reshape(b * f, c, hh, ww).float() if method == "bislerp": diff --git a/modules/sampling/two_stage/tests/test_two_stage.py b/modules/sampling/two_stage/tests/test_two_stage.py index 34b0e14f..ca4a2cb5 100644 --- a/modules/sampling/two_stage/tests/test_two_stage.py +++ b/modules/sampling/two_stage/tests/test_two_stage.py @@ -174,7 +174,7 @@ def test_the_upscaler_list_is_read_again_when_settings_open(upscaler_file, tmp_p assert "later.safetensors" in now -@pytest.mark.parametrize("method", ["bicubic", "mitchell", "catmull-rom", "lanczos", "bislerp", "bilinear", "trilinear", "area", "nearest-exact"]) +@pytest.mark.parametrize("method", ["bicubic", "mitchell", "catmull-rom", "lanczos", "bislerp", "bilinear", "area", "nearest-exact"]) def test_every_plain_method_reaches_the_full_grid_and_keeps_dtype(method): from modules.sampling.two_stage import stages video = torch.randn(1, 24, 2, 5, 7).to(torch.bfloat16)