From a80e58490c2eefe9f828a1f55e27339dd22b08b6 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 10:05:19 +0000 Subject: [PATCH] Two stages as a sampling setting (a model wrapper, any sampler) instead of a preset Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01Tw9fjW75aYB8UyVRZo6fCr --- CHANGELOG.md | 2 +- README.md | 2 +- 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 --------- modules/sampling/two_stage/__init__.py | 54 ++++++++++ modules/sampling/two_stage/stages.py | 101 ++++++++++++++++++ .../two_stage/tests/test_two_stage.py | 68 ++++++++++++ 13 files changed, 236 insertions(+), 226 deletions(-) delete mode 100644 modules/latent/second_pass/tests/test_stage_handoff.py delete mode 100644 modules/sampling/sampler/tests/test_two_stage.py create mode 100644 modules/sampling/two_stage/__init__.py create mode 100644 modules/sampling/two_stage/stages.py create mode 100644 modules/sampling/two_stage/tests/test_two_stage.py diff --git a/CHANGELOG.md b/CHANGELOG.md index e1a5500d..a18d7a44 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,7 @@ ### 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. +- **Two stages** (Settings > Engine > Sampling, H3): the first steps of a run go 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. Works with any sampler; skipped for a run that starts from a picture or a mask. Untested on a GPU. ## [5.1.0] "Classy Chopper" - 2026-10-27 diff --git a/README.md b/README.md index 04c93198..a34e24e4 100644 --- a/README.md +++ b/README.md @@ -166,7 +166,7 @@ I express my deepest gratitude to: - [Inter](https://github.com/rsms/inter) and IBM Plex Mono — the fonts of the app (SIL Open Font License; copies in `app/assets/fonts`); - [musubi-tuner](https://github.com/kohya-ss/musubi-tuner) by kohya-ss (Apache-2.0) — the LoRA training add-on's loop (which layers get a LoRA, per-block checkpointing, noise sampling and loss weighting, optimizer and learning rate) and its block swap (a number of transformer blocks kept in system memory) follow its MiniMax H3 training recipe; - 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 authors of Sol-H3 (arXiv 2609.35110, NVlabs/Sana, Apache-2.0) — the MXFP8 loader option (which blocks stay in bf16) and the mixed resolution two-stage pass follow its recipe; +- 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; Without all of you, this project would've been impossible. diff --git a/core/tests/test_shipped_nodes.py b/core/tests/test_shipped_nodes.py index c93ba519..04a938d0 100644 --- a/core/tests/test_shipped_nodes.py +++ b/core/tests/test_shipped_nodes.py @@ -64,8 +64,7 @@ FROZEN = { "shortcuts", "lorebooks", "reference_intro", "chat", ], "FunUIApplyPromptMarkup": ["model", "clip", "positive", "latent", "markup"], - "FunUILatentResample": ["latent", "upscaler", "operation", "scale", "positive", "negative", "target_width", "target_height"], - "FunUIStageSize": ["width", "height", "scale"], + "FunUILatentResample": ["latent", "upscaler", "operation", "scale", "positive", "negative"], "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 0094700a..beac16cd 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, FunUIStageSize +from .nodes import FunUILatentResample, FunUILatentUpscalerLoader ID = "latent_second_pass" TITLE = "Second pass" @@ -27,4 +27,4 @@ STAGE = "latent" CATEGORY = "sampling" STATUS = "proven" -NODES = [FunUILatentUpscalerLoader, FunUILatentResample, FunUIStageSize] +NODES = [FunUILatentUpscalerLoader, FunUILatentResample] diff --git a/modules/latent/second_pass/nodes.py b/modules/latent/second_pass/nodes.py index 774481f3..7a6efe32 100644 --- a/modules/latent/second_pass/nodes.py +++ b/modules/latent/second_pass/nodes.py @@ -102,35 +102,6 @@ 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: @@ -142,8 +113,7 @@ class FunUILatentResample(io.ComfyNode): "keep the conditioning valid at the new size.", inputs=[ io.Latent.Input("latent"), - 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.LatentUpscaleModel.Input("upscaler"), 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, " @@ -151,10 +121,6 @@ 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"), @@ -165,25 +131,14 @@ class FunUILatentResample(io.ComfyNode): ) @classmethod - 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: + def execute(cls, latent, upscaler, operation: str, scale: float, + positive=None, negative=None) -> 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]) - 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) + up = run_upscaler(upscaler, video, float(scale)) # 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 deleted file mode 100644 index aa701de5..00000000 --- a/modules/latent/second_pass/tests/test_stage_handoff.py +++ /dev/null @@ -1,33 +0,0 @@ -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 06897a63..3ff57647 100644 --- a/modules/models/minimax_h3/pipeline.py +++ b/modules/models/minimax_h3/pipeline.py @@ -168,54 +168,6 @@ 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. @@ -269,10 +221,7 @@ def optional_nodes(): def presets(): return [{"id": "minimax_h3_reference_to_video", "title": "MiniMax 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()}] + "slots": h3_reference_to_video()}] # 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 9f5cf2c1..a7dcdd30 100644 --- a/modules/models/minimax_h3/tests/test_pipeline_preset.py +++ b/modules/models/minimax_h3/tests/test_pipeline_preset.py @@ -33,8 +33,7 @@ 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", - "minimax_h3_reference_to_video_two_stage"] # the second-pass preset is held back + assert [p["id"] for p in offered] == ["minimax_h3_reference_to_video"] # the second-pass preset is held back assert all(p["slots"] for p in offered) @@ -155,23 +154,3 @@ 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 a2e206b1..0806d4dd 100644 --- a/modules/sampling/sampler/nodes.py +++ b/modules/sampling/sampler/nodes.py @@ -19,7 +19,6 @@ 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 @@ -125,7 +124,6 @@ class FunUISampler(io.ComfyNode): outputs=[ io.Latent.Output(display_name="latent"), io.String.Output(display_name="status"), - io.Latent.Output(display_name="denoised"), ], ) @@ -190,8 +188,7 @@ 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. - 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) + callback = latent_preview.prepare_callback(model, max(0, len(sigmas) - 1)) try: samples = guider.sample( @@ -219,20 +216,12 @@ 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]), denoised) + return io.NodeOutput(out, "\n".join([headline, *notes])) # --- 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 deleted file mode 100644 index 02d3967a..00000000 --- a/modules/sampling/sampler/tests/test_two_stage.py +++ /dev/null @@ -1,51 +0,0 @@ -"""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() diff --git a/modules/sampling/two_stage/__init__.py b/modules/sampling/two_stage/__init__.py new file mode 100644 index 00000000..5df64895 --- /dev/null +++ b/modules/sampling/two_stage/__init__.py @@ -0,0 +1,54 @@ +"""Two stages: the first steps of a run on a smaller picture, the last steps at the full size. + +The shot's layout and motion are decided in the first steps, when the picture is still mostly noise, and +a smaller picture decides them nearly as well for a fraction of the work. So the schedule is cut in two: +the first steps run on a picture scaled down, the model's clean guess at that point is resized to the full +size, and the remaining steps finish there. References and text are untouched; only the size of what is +sampled differs. This is the cross-resolution schedule of NVIDIA's Sol-H3 (arXiv 2609.35110); see the +README Thank you section. + +It installs on the model, so it works with any sampler. Not yet run on a real H3 here. +""" + +from comfy.patcher_extension import WrappersMP + +from ..._core import log + +from .stages import make_wrapper + +ID = "two_stage" +TITLE = "Two stages" +MOUNT = "generation.sampling" +STAGE = "load" +CATEGORY = "sampling" +STATUS = "experimental" +REQUIRES = ["audio_stream", "adaln_modalities"] + +SETTINGS = { + "enabled": { + "type": "bool", "default": False, + "label": "Two stages", + "hint": "Runs the first steps on a smaller picture and the last steps at the full size. Cheaper steps " + "up front, the same final size. Skipped for a run that starts from a picture or a mask.", + }, + "size": { + "type": "float", "default": 0.7, "min": 0.4, "max": 0.95, "step": 0.05, + "label": "First stage size", "ui": "slider", + "hint": "Share of the full width and height. 0.7 is about half the work per step.", + "when": {"enabled": True}, + }, + "steps": { + "type": "int", "default": 2, "min": 1, "max": 8, "step": 1, + "label": "Steps in the first stage", + "hint": "Out of the run's own steps. At least one step is always left for the full size.", + "when": {"enabled": True}, + }, +} + + +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)) + return f"first {steps} step(s) at {size:.0%} of the size" diff --git a/modules/sampling/two_stage/stages.py b/modules/sampling/two_stage/stages.py new file mode 100644 index 00000000..5892fd4f --- /dev/null +++ b/modules/sampling/two_stage/stages.py @@ -0,0 +1,101 @@ +"""The wrapper around one sampling run that cuts it in two.""" + +import torch + +from ..._core import log + +TITLE = "FunUI Two stages" + + +def small_grid(h, w, size): + """The first stage's latent grid: scaled, kept even (H3 patches 2x2) and never larger than the full one.""" + def side(n): + return min(n, max(2, int(round(n * size / 2.0)) * 2)) + return side(h), side(w) + + +def _why_not(guider, latent_image, denoise_mask, sigmas, latent_shapes, steps): + if denoise_mask is not None: + return "the run has a mask" + if latent_shapes is None or len(latent_shapes) != 2 or len(latent_shapes[0]) != 5: + return "the latent is not a video with a soundtrack" + if int(torch.count_nonzero(latent_image)) > 0: + return "the run starts from a picture, not from noise" + if len(sigmas) - 1 < 2 or steps >= len(sigmas) - 1: + return f"the run has {len(sigmas) - 1} step(s), too few to cut after {steps}" + for conds in guider.original_conds.values(): + for c in conds: + if "minimax_keyframes" in c.get("model_conds", {}) or "minimax_keyframes" in c: + return "the shot has first/last picture pins sized to the full picture" + try: + from ..ptd.controller import CONTROLLER_KEY + if CONTROLLER_KEY in guider.model_options or CONTROLLER_KEY in guider.model_patcher.model_options: + return "PTD sets its own schedule" + except Exception: # noqa: BLE001 -- no PTD module, nothing to clash with + pass + return None + + +def make_wrapper(size, steps): + def outer(executor, *args, **kwargs): + 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") + guider = executor.class_obj + why = _why_not(guider, latent_image, denoise_mask, sigmas, latent_shapes, steps) + if why is None: + _t, _c, frames, h, w = latent_shapes[0] + h1, w1 = small_grid(h, w, size) + if (h1, w1) == (h, w): + why = "the first stage would be as large as the full picture" + if why is not None: + log.once(f"two_stage:{why}", log.INFO, TITLE, f"not used this time: {why}") + return executor(*args, **kwargs) + + video_noise, audio_noise = comfy.utils.unpack_latents(noise, latent_shapes) + audio_zero = torch.zeros(latent_shapes[1], dtype=latent_image.dtype, device=latent_image.device) + # A crop of the full noise keeps its statistics (a resize would smooth them) and is fixed by the seed. + small_noise, shapes1 = comfy.utils.pack_latents( + (video_noise[..., :h1, :w1].contiguous(), audio_noise)) + small_image, _ = comfy.utils.pack_latents( + (torch.zeros(*latent_shapes[0][:3], h1, w1, dtype=latent_image.dtype, device=latent_image.device), audio_zero)) + first, second = sigmas[:steps + 1], sigmas[steps:] + raw_conds = {k: [c.copy() for c in v] for k, v in guider.conds.items()} + + seen = {} + preview = None + try: + preview = latent_preview.prepare_callback(guider.model_patcher, steps) + except Exception: # noqa: BLE001 -- only the progress bar and previews are lost + pass + + def first_stage_callback(step, x0, x, total): + seen["x0"] = x0 + if preview is not None: + nested = lambda t: comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(t, shapes1)) + preview(step, nested(x0), nested(x), total) + + 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") + executor(small_noise, small_image, sampler, first, None, first_stage_callback, disable_pbar, seed, + latent_shapes=shapes1) + if "x0" not in seen: + raise RuntimeError("the first stage gave no clean guess to hand over") + + video, audio = comfy.utils.unpack_latents(seen["x0"], shapes1) + 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))) + + guider.conds = raw_conds # the first stage encoded the originals for its own size + shifted = None + if callback is not None: + def shifted(step, x0, x, total): + return callback(step + steps, x0, x, total + steps) + return executor(noise, handed, sampler, second, None, shifted, disable_pbar, seed, latent_shapes=latent_shapes) + + return outer diff --git a/modules/sampling/two_stage/tests/test_two_stage.py b/modules/sampling/two_stage/tests/test_two_stage.py new file mode 100644 index 00000000..a8973eb1 --- /dev/null +++ b/modules/sampling/two_stage/tests/test_two_stage.py @@ -0,0 +1,68 @@ +"""Two stages on ComfyUI's real MiniMaxH3 class (a few thousand parameters, CPU), through ComfyUI's own guider: +the wrapper sits on the model, so no FunUI sampler is needed.""" + +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 _run(model, values, sigmas=(1.0, 0.9, 0.7, 0.3, 0.0), latent=None): + import comfy.samplers + from comfy.nested_tensor import NestedTensor + from modules.sampling.two_stage import install + patched = model.clone() + note = install(patched, values, "funui.two_stage") + calls = [] + hook = patched.model.diffusion_model.register_forward_pre_hook( + lambda m, a: calls.append(tuple(a[0][0].shape[-2:]) if isinstance(a[0], (list, tuple)) else tuple(a[0].shape[-2:]))) + guider = comfy.samplers.CFGGuider(patched) + cond = [[torch.randn(1, 5, 48), {}]] + guider.set_conds(cond, cond) + guider.set_cfg(1.0) + latent = latent or NestedTensor((torch.zeros(1, 24, 2, 8, 8), torch.zeros(1, 32, 2, 3))) + noise = NestedTensor((torch.randn(1, 24, 2, 8, 8), torch.randn(1, 32, 2, 3))) + try: + out = guider.sample(noise, latent, comfy.samplers.sampler_object("euler"), torch.tensor(sigmas), + seed=1, disable_pbar=True) + finally: + hook.remove() + return note, out, calls + + +def test_first_steps_run_small_and_the_rest_full_size(model): + note, out, calls = _run(model, {"enabled": True, "size": 0.5, "steps": 2}) + assert "first 2 step(s)" in note + assert calls == [(4, 4), (4, 4), (8, 8), (8, 8)] # 2 steps at half the latent, then 2 at full + video, audio = out.unbind() + assert video.shape == (1, 24, 2, 8, 8) and audio.shape == (1, 32, 2, 3) + assert torch.isfinite(video).all() and torch.isfinite(audio).all() + + +def test_off_means_the_plain_run(model): + note, out, calls = _run(model, {"enabled": False}) + assert note is None and calls == [(8, 8)] * 4 + + +def test_a_run_that_cannot_be_cut_is_left_alone(model): + from comfy.nested_tensor import NestedTensor + started = NestedTensor((torch.ones(1, 24, 2, 8, 8), torch.zeros(1, 32, 2, 3))) + _n, _o, calls = _run(model, {"enabled": True, "size": 0.5, "steps": 2}, latent=started) + assert calls == [(8, 8)] * 4 # a start from a picture is not cut + _n, _o, calls = _run(model, {"enabled": True, "size": 0.5, "steps": 4}) + assert calls == [(8, 8)] * 4 # nothing would be left for the full size