From 791bb6ca484f6d32ce4b69da59f4f8f81eecf576 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 09:55:37 +0000 Subject: [PATCH] H3: two-stage preset (first steps on a smaller picture), Sampler denoised output, Resample to an exact size without an upscaler file, Stage Size node Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01Tw9fjW75aYB8UyVRZo6fCr --- CHANGELOG.md | 7 +++ core/tests/test_shipped_nodes.py | 3 +- modules/latent/second_pass/__init__.py | 4 +- modules/latent/second_pass/nodes.py | 53 +++++++++++++++++-- .../second_pass/tests/test_stage_handoff.py | 33 ++++++++++++ modules/models/minimax_h3/pipeline.py | 53 ++++++++++++++++++- .../minimax_h3/tests/test_pipeline_preset.py | 23 +++++++- modules/sampling/sampler/nodes.py | 15 +++++- .../sampling/sampler/tests/test_two_stage.py | 51 ++++++++++++++++++ 9 files changed, 231 insertions(+), 11 deletions(-) create mode 100644 modules/latent/second_pass/tests/test_stage_handoff.py create mode 100644 modules/sampling/sampler/tests/test_two_stage.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 82af093d..e1a5500d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,12 @@ # Changelog +## [Unreleased] + +### Added + +- **H3 loader: MXFP8** (Blackwell GPUs): stores the attention and MLP weights of the middle blocks as MXFP8 (first 2 and last 3 blocks stay bf16) for about twice the matmul rate. Needs torch 2.10+ with CUDA 13; skipped with a note elsewhere. Untested on a GPU; listen to the sound first. +- **H3 preset "Reference to Video, two stages (test)"**: the model's own schedule is cut in two. The first steps run on a smaller picture (First stage size, 0.7 by default), the model's clean guess is resized to the full size, and the last steps finish there. New nodes: Stage Size; Sampler gained a `denoised` output; Latent Resample can land on an exact size and works without an upscaler file (bicubic). Untested on a GPU. + ## [5.1.0] "Classy Chopper" - 2026-10-27 ### Changed diff --git a/core/tests/test_shipped_nodes.py b/core/tests/test_shipped_nodes.py index 04a938d0..c93ba519 100644 --- a/core/tests/test_shipped_nodes.py +++ b/core/tests/test_shipped_nodes.py @@ -64,7 +64,8 @@ FROZEN = { "shortcuts", "lorebooks", "reference_intro", "chat", ], "FunUIApplyPromptMarkup": ["model", "clip", "positive", "latent", "markup"], - "FunUILatentResample": ["latent", "upscaler", "operation", "scale", "positive", "negative"], + "FunUILatentResample": ["latent", "upscaler", "operation", "scale", "positive", "negative", "target_width", "target_height"], + "FunUIStageSize": ["width", "height", "scale"], "FunUIModifierSettings": ["settings"], "FunUIGuideFrames": ["positive", "negative", "vae", "latent", "frames"], "FunUIIdentityTransfer": ["model", "vae", "positive", "negative", "latent", "image", "identity_projector", diff --git a/modules/latent/second_pass/__init__.py b/modules/latent/second_pass/__init__.py index beac16cd..0094700a 100644 --- a/modules/latent/second_pass/__init__.py +++ b/modules/latent/second_pass/__init__.py @@ -19,7 +19,7 @@ What this module adds is the part the graph could not already do: v4's user verdict: the H3 second pass WORKS (2026-09-20). """ -from .nodes import FunUILatentResample, FunUILatentUpscalerLoader +from .nodes import FunUILatentResample, FunUILatentUpscalerLoader, FunUIStageSize ID = "latent_second_pass" TITLE = "Second pass" @@ -27,4 +27,4 @@ STAGE = "latent" CATEGORY = "sampling" STATUS = "proven" -NODES = [FunUILatentUpscalerLoader, FunUILatentResample] +NODES = [FunUILatentUpscalerLoader, FunUILatentResample, FunUIStageSize] diff --git a/modules/latent/second_pass/nodes.py b/modules/latent/second_pass/nodes.py index 7a6efe32..774481f3 100644 --- a/modules/latent/second_pass/nodes.py +++ b/modules/latent/second_pass/nodes.py @@ -102,6 +102,35 @@ def rescale(conditioning, h, w): return conditioning, 0 +def bicubic_to(video, h, w): + """The plain way up when there is no upscaler file: bicubic per frame, same as the way down.""" + return downscale_to(video, h, w) + + +class FunUIStageSize(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="FunUIStageSize", + display_name="FunUI Stage Size", + category="FunUI/Latent", + description="The size of a cheaper first stage: the full size scaled down, kept on the 32 pixel grid.", + inputs=[ + io.Int.Input("width", default=640, min=32, max=16384, step=32), + io.Int.Input("height", default=640, min=32, max=16384, step=32), + io.Float.Input("scale", default=0.7, min=0.25, max=1.0, step=0.05, + tooltip="Share of the full width and height the first stage runs at. 0.7 is about half the work per step."), + ], + outputs=[io.Int.Output(display_name="width"), io.Int.Output(display_name="height")], + ) + + @classmethod + def execute(cls, width: int, height: int, scale: float) -> io.NodeOutput: + def snap(v): + return max(64, int(round(v * scale / 32.0)) * 32) + return io.NodeOutput(min(snap(width), int(width)), min(snap(height), int(height))) + + class FunUILatentResample(io.ComfyNode): @classmethod def define_schema(cls) -> io.Schema: @@ -113,7 +142,8 @@ class FunUILatentResample(io.ComfyNode): "keep the conditioning valid at the new size.", inputs=[ io.Latent.Input("latent"), - io.LatentUpscaleModel.Input("upscaler"), + io.LatentUpscaleModel.Input("upscaler", optional=True, + tooltip="Without one the latent is resized with bicubic, which is enough to hand a first stage to a second."), io.Combo.Input("operation", options=["sharpen", "upscale"], default="sharpen", tooltip="sharpen: up, then back to this size -- pass 2 " "re-denoises the added detail. upscale: kept larger, " @@ -121,6 +151,10 @@ class FunUILatentResample(io.ComfyNode): io.Float.Input("scale", default=2.0, min=1.0, max=4.0, step=0.05), io.Conditioning.Input("positive", optional=True), io.Conditioning.Input("negative", optional=True), + io.Int.Input("target_width", default=0, min=0, max=16384, step=32, optional=True, + tooltip="upscale only: land exactly on this picture width (0 keeps whatever scale gives)."), + io.Int.Input("target_height", default=0, min=0, max=16384, step=32, optional=True, + tooltip="upscale only: land exactly on this picture height."), ], outputs=[ io.Latent.Output(display_name="latent"), @@ -131,14 +165,25 @@ class FunUILatentResample(io.ComfyNode): ) @classmethod - def execute(cls, latent, upscaler, operation: str, scale: float, - positive=None, negative=None) -> io.NodeOutput: + def execute(cls, latent, upscaler=None, operation: str = "sharpen", scale: float = 2.0, + positive=None, negative=None, target_width: int = 0, target_height: int = 0) -> io.NodeOutput: video, rebuild = _video_of(latent["samples"]) if getattr(video, "ndim", 0) != 5: raise RuntimeError( f"expected a video latent [B, C, T, H, W], got {tuple(video.shape)}") h, w = int(video.shape[-2]), int(video.shape[-1]) - up = run_upscaler(upscaler, video, float(scale)) + goal = None + if operation == "upscale" and target_width and target_height: + goal = (int(target_height) // 16, int(target_width) // 16) # H3's picture is 16 pixels per latent cell + scale = max(1.0, goal[0] / h, goal[1] / w) + if upscaler is None: + if operation == "sharpen": + raise RuntimeError("sharpen needs a latent upscaler: connect one, or use upscale.") + up = bicubic_to(video, max(2, int(round(h * scale / 2.0)) * 2), max(2, int(round(w * scale / 2.0)) * 2)) + else: + up = run_upscaler(upscaler, video, float(scale)) + if goal is not None: + up = downscale_to(up, *goal) if up.shape[-2] >= goal[0] and up.shape[-1] >= goal[1] else bicubic_to(up, *goal) # noise_mask is kept: a pinned frame must stay pinned in pass 2, and # ComfyUI's prepare_mask resizes a mask to whatever latent it meets. diff --git a/modules/latent/second_pass/tests/test_stage_handoff.py b/modules/latent/second_pass/tests/test_stage_handoff.py new file mode 100644 index 00000000..aa701de5 --- /dev/null +++ b/modules/latent/second_pass/tests/test_stage_handoff.py @@ -0,0 +1,33 @@ +import asyncio + +import torch + + +def _run(coro_or_value): + return coro_or_value + + +def test_stage_size_snaps_to_the_grid_and_never_grows(): + from ..nodes import FunUIStageSize + w, h = FunUIStageSize.execute(640, 352, 0.7) + assert (w, h) == (448, 256) + assert all(v % 32 == 0 for v in (w, h)) + assert tuple(FunUIStageSize.execute(96, 96, 1.0)) == (96, 96) + assert tuple(FunUIStageSize.execute(64, 64, 0.25)) == (64, 64) # never below 64 + + +def test_resample_lands_on_the_target_grid_without_an_upscaler_file(): + from ..nodes import FunUILatentResample + video = torch.randn(1, 24, 5, 16, 28) # 256x448 pixels + out = FunUILatentResample.execute({"samples": video}, None, "upscale", 1.5, None, None, 640, 352) + assert out.result[0]["samples"].shape == (1, 24, 5, 22, 40) # 352x640 + assert "upscale" in out.result[3] + + +def test_resample_hands_the_audio_part_through_untouched(): + from comfy.nested_tensor import NestedTensor + from ..nodes import FunUILatentResample + video, audio = torch.randn(1, 24, 5, 8, 8), torch.randn(1, 32, 2, 10) + out = FunUILatentResample.execute({"samples": NestedTensor((video, audio))}, None, "upscale", 2.0, None, None, 256, 256) + v, a = out.result[0]["samples"].unbind() + assert v.shape == (1, 24, 5, 16, 16) and torch.equal(a, audio) diff --git a/modules/models/minimax_h3/pipeline.py b/modules/models/minimax_h3/pipeline.py index 3ff57647..06897a63 100644 --- a/modules/models/minimax_h3/pipeline.py +++ b/modules/models/minimax_h3/pipeline.py @@ -168,6 +168,54 @@ def h3_reference_to_video(): return slots +def h3_reference_to_video_two_stage(): + """The same pipeline run in two stages of different size (the cross-resolution schedule of NVIDIA's Sol-H3). + + The model's own schedule is cut in two. The first steps (where the shot's layout and motion are + decided) run on a smaller picture, the model's clean guess at that point is resized to the full + size, and the last steps run there, so most of the early work costs a fraction. References and + text ride along unchanged: only the empty latent differs in size. Untested on a GPU. + """ + slots = h3_reference_to_video() + slots = [s for s in slots if s["id"] != "sampler"] + at = next(i for i, s in enumerate(slots) if s["id"] == "decode") + sampling_roles = [{"at": "generation.sampling", "input": "sampler_name", "label": "Sampler"}, + {"at": "generation.seed", "input": "seed", "label": "Seed"}] + slots[at:at] = [ + {"id": "stage_size", "group": "Sampling", "node": "FunUIStageSize", + "roles": [{"at": "generation.sampling", "input": "scale", "label": "First stage size"}], + "inputs": {"width": ["image_transform", 4], "height": ["image_transform", 5], "scale": 0.7}}, + {"id": "latent_low", "group": "Sampling", "node": "FunUIEmptyLatent", + "roles": [{"at": "project.video", "input": "length", "label": "Length", "drives": "frames"}], + "inputs": {"model": ["model", 0], "width": ["stage_size", 0], "height": ["stage_size", 1], + "length": 124, "batch_size": 1, "audio_vae": ["audio_vae", 0], "frame_rate": 24.0}}, + {"id": "schedule", "group": "Sampling", "node": "BasicScheduler", + "roles": [{"at": "generation.sampling", "input": "steps", "label": "Steps"}, + {"at": "generation.sampling", "input": "scheduler", "label": "Scheduler"}], + "inputs": {"model": ["markup_apply", 0], "scheduler": "normal", "steps": 4, "denoise": 1.0}}, + {"id": "split", "group": "Sampling", "node": "SplitSigmas", + "roles": [{"at": "generation.sampling", "input": "step", "label": "Steps in the first stage"}], + "inputs": {"sigmas": ["schedule", 0], "step": 2}}, + {"id": "sampler_low", "group": "Sampling", "node": "FunUISampler", + "roles": sampling_roles, + "inputs": { + "model": ["markup_apply", 0], "positive": ["erase", 0], "negative": ["negative", 0], + "latent": ["latent_low", 0], "settings": ["settings", 0], "sigmas": ["split", 0], + "seed": 0, "steps": 4, "cfg": 1.0, "sampler_name": "euler", "scheduler": "normal", "denoise": 1.0}}, + {"id": "resample", "group": "Sampling", "node": "FunUILatentResample", + "inputs": {"latent": ["sampler_low", 2], "operation": "upscale", "scale": 1.5, + "positive": ["erase", 0], "negative": ["negative", 0], + "target_width": ["image_transform", 4], "target_height": ["image_transform", 5]}}, + {"id": "sampler", "group": "Sampling", "node": "FunUISampler", + "roles": sampling_roles, + "inputs": { + "model": ["markup_apply", 0], "positive": ["resample", 1], "negative": ["resample", 2], + "latent": ["resample", 0], "settings": ["settings", 0], "sigmas": ["split", 1], + "seed": 0, "steps": 4, "cfg": 1.0, "sampler_name": "euler", "scheduler": "normal", "denoise": 1.0}}, + ] + return slots + + def h3_reference_to_video_second_pass(): """The same pipeline, then a second pass: resample between passes (sharpen by default) and a second sampler running a typed schedule as written. @@ -221,7 +269,10 @@ def optional_nodes(): def presets(): return [{"id": "minimax_h3_reference_to_video", "title": "MiniMax H3 · Reference to Video", - "slots": h3_reference_to_video()}] + "slots": h3_reference_to_video()}, + {"id": "minimax_h3_reference_to_video_two_stage", + "title": "MiniMax H3 · Reference to Video, two stages (test)", + "slots": h3_reference_to_video_two_stage()}] # Held back, not offered: h3_reference_to_video_second_pass() builds clean but was never verified on a GPU # (it needs a compatible H3 latent upscaler file). Add it back here once it has been. diff --git a/modules/models/minimax_h3/tests/test_pipeline_preset.py b/modules/models/minimax_h3/tests/test_pipeline_preset.py index a7dcdd30..9f5cf2c1 100644 --- a/modules/models/minimax_h3/tests/test_pipeline_preset.py +++ b/modules/models/minimax_h3/tests/test_pipeline_preset.py @@ -33,7 +33,8 @@ def registered(comfyui): def test_the_preset_is_offered(registered): from modules.models.minimax_h3.pipeline import presets offered = presets() - assert [p["id"] for p in offered] == ["minimax_h3_reference_to_video"] # the second-pass preset is held back + assert [p["id"] for p in offered] == ["minimax_h3_reference_to_video", + "minimax_h3_reference_to_video_two_stage"] # the second-pass preset is held back assert all(p["slots"] for p in offered) @@ -154,3 +155,23 @@ def test_the_second_pass_preset_builds_clean_and_decodes_pass_two(registered): assert [p for p in problems if not any(m in p for m in unset)] == [] decode = next(s for s in slots if s["id"] == "decode") assert decode["inputs"]["samples"] == ["pass2", 0] + + +def test_the_two_stage_preset_is_wired_to_real_slots_and_cuts_the_schedule_once(): + from modules.models.minimax_h3.pipeline import h3_reference_to_video_two_stage + + slots = h3_reference_to_video_two_stage() + ids = [s["id"] for s in slots] + assert len(ids) == len(set(ids)) + known = set(ids) + by_id = {s["id"]: s for s in slots} + for slot in slots: + for name, value in (slot.get("inputs") or {}).items(): + if isinstance(value, list) and len(value) == 2 and isinstance(value[0], str): + assert value[0] in known, f"{slot['id']}.{name} is wired to {value[0]!r}, which is not a slot here" + # First stage takes the front of the schedule, the second stage the rest of the same schedule. + assert by_id["sampler_low"]["inputs"]["sigmas"] == ["split", 0] and by_id["sampler"]["inputs"]["sigmas"] == ["split", 1] + # The second stage starts from the first stage's clean guess (output 2), resized, and decode reads the second stage. + assert by_id["resample"]["inputs"]["latent"] == ["sampler_low", 2] + assert by_id["sampler"]["inputs"]["latent"] == ["resample", 0] + assert by_id["decode"]["inputs"]["samples"] == ["sampler", 0] diff --git a/modules/sampling/sampler/nodes.py b/modules/sampling/sampler/nodes.py index 0806d4dd..a2e206b1 100644 --- a/modules/sampling/sampler/nodes.py +++ b/modules/sampling/sampler/nodes.py @@ -19,6 +19,7 @@ Three things keep it from becoming a hub again: """ import comfy.model_management +import comfy.nested_tensor import comfy.sample import comfy.samplers import comfy.utils @@ -124,6 +125,7 @@ class FunUISampler(io.ComfyNode): outputs=[ io.Latent.Output(display_name="latent"), io.String.Output(display_name="status"), + io.Latent.Output(display_name="denoised"), ], ) @@ -188,7 +190,8 @@ class FunUISampler(io.ComfyNode): # Progress and previews, from ComfyUI's own machinery rather than ours: # it already picks a preview method, already falls back when the decoder # is missing, and already reports steps to whatever is watching. - callback = latent_preview.prepare_callback(model, max(0, len(sigmas) - 1)) + x0_output = {} # the model's last clean guess: what a run cut short at a middle sigma hands to the next stage + callback = latent_preview.prepare_callback(model, max(0, len(sigmas) - 1), x0_output) try: samples = guider.sample( @@ -216,12 +219,20 @@ class FunUISampler(io.ComfyNode): out.pop("downscale_ratio_spacial", None) out.pop("downscale_ratio_temporal", None) out["samples"] = samples + denoised = out + if "x0" in x0_output: + x0 = x0_output["x0"] + if samples.is_nested and not x0.is_nested: + x0 = comfy.nested_tensor.NestedTensor( + comfy.utils.unpack_latents(x0, [t.shape for t in samples.unbind()])) + denoised = out.copy() + denoised["samples"] = model.model.process_latent_out(x0.cpu()).to(comfy.model_management.intermediate_device()) schedule = f"wired schedule from {float(sigmas[0]):.3g}" if wired else scheduler headline = (f"{run}: {sampler_name} / {schedule}, {len(sigmas) - 1} step(s), " f"cfg {cfg}") log.info("FunUI Sampler", headline) - return io.NodeOutput(out, "\n".join([headline, *notes])) + return io.NodeOutput(out, "\n".join([headline, *notes]), denoised) # --- the parts, kept separate so each is testable on its own ----------- diff --git a/modules/sampling/sampler/tests/test_two_stage.py b/modules/sampling/sampler/tests/test_two_stage.py new file mode 100644 index 00000000..02d3967a --- /dev/null +++ b/modules/sampling/sampler/tests/test_two_stage.py @@ -0,0 +1,51 @@ +"""The two-stage hand-off on ComfyUI's real MiniMaxH3 class (a few thousand parameters, CPU): a first stage cut +short at a middle sigma on a small picture, its clean guess resized, a second stage finishing the same +schedule at the full size.""" + +import pytest +import torch + + +@pytest.fixture(scope="module") +def model(comfyui): + import comfy.model_patcher as mp + import comfy.supported_models as sm + torch.manual_seed(0) + config = sm.MiniMaxH3(dict( + image_model="minimax_h3", hidden_size=64, num_layers=4, token_refiner_num_layers=1, num_attention_heads=2, + attention_head_dim=32, ffn_hidden_size=64, text_dim=48, timestep_input_dim=32, time_embed_hidden_size=64, + time_embed_dim=32, rope_inv_freq_len=4)) + config.set_inference_dtype(torch.float32, None) + base = config.get_model({}, device=torch.device("cpu")) + for p in base.diffusion_model.parameters(): + torch.nn.init.normal_(p, std=0.05) + p.requires_grad_(False) + return mp.ModelPatcher(base, load_device=torch.device("cpu"), offload_device=torch.device("cpu")) + + +def test_a_run_cut_short_hands_over_a_clean_guess_and_a_second_stage_finishes_it_bigger(model): + from comfy.nested_tensor import NestedTensor + from modules.latent.second_pass.nodes import FunUILatentResample + from modules.sampling.sampler.nodes import FunUISampler + + cond = [[torch.randn(1, 5, 48), {}]] + small = {"samples": NestedTensor((torch.zeros(1, 24, 2, 4, 4), torch.zeros(1, 32, 2, 3)))} + sigmas = torch.tensor([1.0, 0.9, 0.7, 0.3, 0.0]) + first = FunUISampler.execute(model, cond, cond, small, seed=1, steps=4, cfg=1.0, sampler_name="euler", + scheduler="normal", sigmas=sigmas[:3]) + latent, _status, denoised = first.result + assert len(first.result) == 3 + v_noisy, a_noisy = latent["samples"].unbind() + v_clean, a_clean = denoised["samples"].unbind() + assert v_clean.shape == v_noisy.shape and a_clean.shape == a_noisy.shape + assert all(torch.isfinite(t).all() for t in (v_noisy, a_noisy, v_clean, a_clean)) + assert not torch.equal(v_clean, v_noisy) # cut at sigma 0.7: the guess is not the noisy state + + big = FunUILatentResample.execute(denoised, None, "upscale", 2.0, cond, cond, 128, 128).result + assert big[0]["samples"].unbind()[0].shape == (1, 24, 2, 8, 8) + + second = FunUISampler.execute(model, big[1], big[2], big[0], seed=1, steps=4, cfg=1.0, sampler_name="euler", + scheduler="normal", sigmas=sigmas[2:]) + v, a = second.result[0]["samples"].unbind() + assert v.shape == (1, 24, 2, 8, 8) and a.shape == a_noisy.shape + assert torch.isfinite(v).all() and torch.isfinite(a).all()