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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Tw9fjW75aYB8UyVRZo6fCr

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

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",

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]

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.

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)

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.
 

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]

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 -----------
 

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()