Two stages as a sampling setting (a model wrapper, any sampler) instead of a preset
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Tw9fjW75aYB8UyVRZo6fCr
This commit is contained in:
Claude
+1
-1
@@ -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
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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 -----------
|
||||
|
||||
|
||||
@@ -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()
|
||||
@@ -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"
|
||||
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user