One node per capability and no mode dropdown, because a mode widget spanning two sampling paths is a mistake this line has already paid for. Shared plumbing sits at module scope so a second node cannot drift from this one. Built on ComfyUI core alone — no second pack to install, which is the whole premise. Core's own fl2va node supplies the conditioning and keyframes; the sampling and the two decodes mirror core's SamplerCustomAdvanced, VAEDecode and vae_decode_audio, and say so where they do, so they can be diffed against the source rather than trusted. What it refuses to offer: sampler, scheduler, cfg, sigma shift, negative prompt. What it guards: the checkpoint family (fl2va vs ref2va, read out of the workflow because nothing in the MODEL object carries the filename), the canvas — core does NOT round the output canvas on this path, so the menu is the only guard — and both frames' fitting, which core does asymmetrically and silently. 66 offline checks in tests/test_offline.py, no pytest and no GPU needed. They caught the canvas claim above, which an earlier draft had backwards.
419 lines
18 KiB
Python
419 lines
18 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""Shared plumbing for every SimpleShot node.
|
|
|
|
Written once and used by all of them, so a second node cannot quietly drift
|
|
from the first: the duration and canvas tables, the family guard, the sampling
|
|
call and the two decodes all live here. Nodes above this file choose WHAT to
|
|
offer; this file decides HOW it runs.
|
|
|
|
Everything here talks to ComfyUI core only — no second custom pack, by design.
|
|
The functions that mirror a core node's body say which node they mirror, so the
|
|
next person can diff them against the source rather than trust a comment.
|
|
"""
|
|
# ── the two tables the user actually chooses from ────────────────────────────
|
|
#
|
|
# THE TWO ARE NOT ALIKE, and an early version of this file had it wrong.
|
|
#
|
|
# Frame count: core DOES correct it. `temporal_shape` runs `align_frame_count`
|
|
# on the way into the latent, so any length is snapped up to the 17k+5 grid.
|
|
# The entries here are already on the grid, so nothing is corrected and nothing
|
|
# needs to be. The menu is convenience.
|
|
#
|
|
# Canvas: core does NOT correct it. `adapt_canvas` exists in the same module and
|
|
# looks like it would, but it is only ever called for reference VIDEOS in the
|
|
# ref2va node — the first/last-frame path passes width and height straight into
|
|
# `_empty_av_latent`, which divides them by 16. So an off-multiple canvas is a
|
|
# broken tensor, not a corrected one, and THIS MENU IS THE GUARD. Every entry
|
|
# is a multiple of 32 and stays inside the 768*1344 area the model works in.
|
|
# Read the source before adding one; tests/test_offline.py enforces both rules.
|
|
|
|
DURATIONS = {
|
|
"2 s — quick test (56 frames)": 56,
|
|
"5 s (124 frames)": 124,
|
|
"10 s (243 frames)": 243,
|
|
"15 s (362 frames)": 362,
|
|
}
|
|
DURATION_NAMES = list(DURATIONS)
|
|
DURATION_DEFAULT = "5 s (124 frames)"
|
|
|
|
RESOLUTIONS = {
|
|
"portrait 768 x 1344": (768, 1344),
|
|
"landscape 1344 x 768": (1344, 768),
|
|
"portrait 736 x 1280 - cheaper": (736, 1280),
|
|
"landscape 1280 x 736 - cheaper": (1280, 736),
|
|
"square 960 x 960": (960, 960),
|
|
}
|
|
RESOLUTION_NAMES = list(RESOLUTIONS)
|
|
RESOLUTION_DEFAULT = "portrait 768 x 1344" # core's own default, and the one
|
|
# every render behind this pack used
|
|
|
|
# ── the sampling pair, deliberately not exposed ──────────────────────────────
|
|
#
|
|
# This is the pairing ComfyUI's own MiniMax-H3 template ships (res_multistep +
|
|
# simple, cfg-free guidance, no sigma shift), which is the closest thing to an
|
|
# author-shipped config for a core-only fl2va graph. It is a constant and not a
|
|
# widget on purpose: exposing a sampler only lets a working thing be broken, and
|
|
# this line sells "it works the first time". One edit here changes every node.
|
|
SAMPLER_NAME = "res_multistep"
|
|
SCHEDULER = "simple"
|
|
STEPS_DEFAULT = 20
|
|
|
|
FPS = 24 # H3 renders at 24 fps. Not configurable — see the spec.
|
|
|
|
|
|
# ── core, imported late and with a legible failure ───────────────────────────
|
|
def _fail(msg):
|
|
"""One exception type, one voice. Errors are ours, not a stack trace."""
|
|
raise RuntimeError("[SimpleShot] " + msg)
|
|
|
|
|
|
def core_h3():
|
|
"""core's MiniMax-H3 module, or a sentence saying what to do about it."""
|
|
try:
|
|
import comfy_extras.nodes_minimax_h3 as h3
|
|
except Exception:
|
|
_fail("this ComfyUI has no MiniMax-H3 support built in, so the model "
|
|
"cannot run at all. Update ComfyUI to a version that ships "
|
|
"comfy_extras/nodes_minimax_h3.py (Manager > Update ComfyUI, "
|
|
"then restart). Nothing else needs installing.")
|
|
if not hasattr(h3, "MiniMaxH3ImageToVideo"):
|
|
_fail("this ComfyUI ships MiniMax-H3 support, but not the "
|
|
"first/last-frame node this workflow is built on. Update "
|
|
"ComfyUI and restart.")
|
|
return h3
|
|
|
|
|
|
def unwrap(out):
|
|
"""The values out of a core node's return, old API or new.
|
|
|
|
Core's newer nodes return a NodeOutput (indexable, with `.result`); the
|
|
older ones return a plain tuple. Both are handled so this pack does not
|
|
break on the version boundary.
|
|
"""
|
|
result = getattr(out, "result", None)
|
|
if result is not None:
|
|
return tuple(result)
|
|
if isinstance(out, (tuple, list)):
|
|
return tuple(out)
|
|
return (out,)
|
|
|
|
|
|
# ── what the model will actually do with the numbers ─────────────────────────
|
|
def align_frame_count(n):
|
|
"""Mirror of core's `align_frame_count`: snap UP to the 17k+5 grid."""
|
|
n = max(5, int(n))
|
|
while n % 17 != 5:
|
|
n += 1
|
|
return n
|
|
|
|
|
|
MAX_PIXELS = 768 * 1344 # the area core works in
|
|
CANVAS_MULTIPLE = 32 # the latent divides by 16; 32 keeps the video vae happy
|
|
|
|
|
|
def canvas_is_legal(width, height):
|
|
"""(ok, reason) for a canvas, since core will not fix a bad one.
|
|
|
|
Not a mirror of core's `adapt_canvas` — that function is for reference
|
|
videos and is never applied to the output canvas on this path. This is the
|
|
check core does not do.
|
|
"""
|
|
if width % CANVAS_MULTIPLE or height % CANVAS_MULTIPLE:
|
|
return (False, "%dx%d is not a multiple of %d, and core does not round "
|
|
"the output canvas — the latent would be the wrong shape"
|
|
% (width, height, CANVAS_MULTIPLE))
|
|
if width * height > MAX_PIXELS:
|
|
return (False, "%dx%d is %.2f MP, past the %.2f MP the model works in"
|
|
% (width, height, width * height / 1e6, MAX_PIXELS / 1e6))
|
|
return (True, "")
|
|
|
|
|
|
def resolve_shape(duration, resolution):
|
|
"""(frames, width, height, notes) for the two dropdown values.
|
|
|
|
Reads the tables, applies core's frame snapping so the node can SAY what the
|
|
model will use rather than what the menu said, and refuses a canvas core
|
|
would not fix. Silence about an adjustment is how a surprise becomes a bug
|
|
report.
|
|
"""
|
|
if duration not in DURATIONS:
|
|
_fail("duration=%r is not one of %s. (A duration arriving down a link "
|
|
"is never checked against the menu by ComfyUI, which is why this "
|
|
"is checked here.)" % (duration, DURATION_NAMES))
|
|
if resolution not in RESOLUTIONS:
|
|
_fail("resolution=%r is not one of %s." % (resolution, RESOLUTION_NAMES))
|
|
|
|
asked_frames = DURATIONS[duration]
|
|
width, height = RESOLUTIONS[resolution]
|
|
frames = align_frame_count(asked_frames)
|
|
|
|
ok, why = canvas_is_legal(width, height)
|
|
if not ok:
|
|
_fail("resolution=%r maps to an illegal canvas: %s. This is a bug in "
|
|
"the pack, not in your graph — please report it." % (resolution, why))
|
|
|
|
notes = []
|
|
if frames != asked_frames:
|
|
notes.append("frame count %d -> %d (core snaps up to its 17k+5 grid)"
|
|
% (asked_frames, frames))
|
|
return frames, width, height, notes
|
|
|
|
|
|
# ── which checkpoint is upstream, and is it the right family ─────────────────
|
|
LOADER_TYPES = (
|
|
"UNETLoader", "UnetLoaderGGUF", "UnetLoaderGGUFAdvanced",
|
|
"CheckpointLoaderSimple", "H3ModelLoaderAny",
|
|
)
|
|
|
|
|
|
def _walk_back(workflow, node_id, input_name, depth=0):
|
|
"""The node feeding `input_name` of `node_id`, through reroutes and bypasses.
|
|
|
|
Returns (type, first_widget_value) or (None, None). Never raises: a guard
|
|
that can crash is worse than a guard that admits it does not know.
|
|
|
|
Reroutes are followed by TAKING THE FIRST INPUT'S LINK rather than by name,
|
|
because a Reroute's input has no name at all — matching on one would quietly
|
|
end the walk at the reroute and turn a real answer into "cannot tell".
|
|
"""
|
|
if depth > 12:
|
|
return (None, None)
|
|
nodes = {int(n["id"]): n for n in workflow.get("nodes", [])}
|
|
node = nodes.get(int(node_id))
|
|
if node is None:
|
|
return (None, None)
|
|
|
|
link_id = None
|
|
for index, slot in enumerate(node.get("inputs") or []):
|
|
if slot.get("name") == input_name or (input_name is None and index == 0):
|
|
link_id = slot.get("link")
|
|
break
|
|
if link_id is None:
|
|
return (None, None)
|
|
|
|
for link in workflow.get("links") or []:
|
|
# [link_id, origin_id, origin_slot, target_id, target_slot, type]
|
|
if not isinstance(link, (list, tuple)) or len(link) < 5:
|
|
continue
|
|
if link[0] != link_id:
|
|
continue
|
|
origin = nodes.get(int(link[1]))
|
|
if origin is None:
|
|
return (None, None)
|
|
kind = str(origin.get("type") or "")
|
|
# A reroute, or a BYPASSED node (mode 4), is not the source: keep going
|
|
# through whatever feeds it.
|
|
if kind in ("Reroute", "RerouteNode", "PrimitiveNode") or origin.get("mode") == 4:
|
|
slots = origin.get("inputs") or []
|
|
passthrough = slots[0].get("name") if slots else None
|
|
return _walk_back(workflow, link[1], passthrough or None, depth + 1)
|
|
widgets = origin.get("widgets_values") or []
|
|
first = widgets[0] if widgets else None
|
|
return (kind, first if isinstance(first, str) else None)
|
|
return (None, None)
|
|
|
|
|
|
def checkpoint_note(model, extra_pnginfo, unique_id, want="fl2va"):
|
|
"""A line about the loaded checkpoint, and a hard stop on the wrong family.
|
|
|
|
fl2va and ref2va are different trainings of the same architecture and load
|
|
through the same loader, so nothing in the MODEL object identifies which one
|
|
arrived. The name does — and it is upstream in the graph, not in the object,
|
|
so it is read out of the workflow the frontend sends along.
|
|
|
|
Deliberately soft when it cannot tell: a run refused on a guess is worse
|
|
than a run that says "I could not check this". Deliberately hard when it
|
|
can: fl2va has no reference rows and ref2va cannot land on a last frame, so
|
|
the wrong file is a wasted render, not a worse one.
|
|
"""
|
|
name = None
|
|
try:
|
|
name = str(getattr(getattr(model, "model", None),
|
|
"h3_checkpoint_name", "") or "") or None
|
|
except Exception:
|
|
name = None
|
|
|
|
if name is None and extra_pnginfo and unique_id is not None:
|
|
try:
|
|
workflow = (extra_pnginfo or {}).get("workflow") or {}
|
|
_kind, value = _walk_back(workflow, unique_id, "model")
|
|
name = value
|
|
except Exception:
|
|
name = None
|
|
|
|
if not name:
|
|
return ("checkpoint: name unavailable, family UNCHECKED — if the render "
|
|
"ignores your frames, confirm you loaded an fl2va file")
|
|
|
|
low = name.lower()
|
|
family = "fl2va" if "fl2va" in low else ("ref2va" if "ref2va" in low else None)
|
|
if family is None:
|
|
return ("checkpoint: %s — family unproven (the filename says neither "
|
|
"fl2va nor ref2va), so this was not checked" % name)
|
|
if family != want:
|
|
_fail("wrong checkpoint for this node.\n"
|
|
" Loaded: %s (a %s model).\n"
|
|
" This node is the FIRST/LAST-FRAME task and needs an %s file.\n"
|
|
" %s renders from reference photographs and cannot be aimed at a "
|
|
"last frame; %s renders between the frames you give it and has no "
|
|
"reference inputs at all.\n"
|
|
" Fix: pick the %s checkpoint in the loader above."
|
|
% (name, family, want, family, want, want))
|
|
return "checkpoint: %s (%s, correct for this node)" % (name, family)
|
|
|
|
|
|
# ── images in ────────────────────────────────────────────────────────────────
|
|
def image_notes(first_frame, last_frame, width, height):
|
|
"""What the two optional images will have done to them, said in advance.
|
|
|
|
Core treats them differently and neither the UI nor the log mentions it:
|
|
first_frame is the geometry anchor and is STRETCHED to the canvas, while
|
|
last_frame is a follower and is CENTRE-CROPPED. A portrait photo dropped
|
|
into a landscape shot therefore distorts in one slot and loses its edges in
|
|
the other — the exact kind of surprise this pack exists to remove.
|
|
"""
|
|
notes = []
|
|
|
|
def describe(img, label, behaviour):
|
|
try:
|
|
h, w = int(img.shape[1]), int(img.shape[2])
|
|
except Exception:
|
|
return
|
|
if (w, h) == (width, height):
|
|
notes.append("%s: %dx%d, matches the canvas" % (label, w, h))
|
|
return
|
|
same_aspect = abs((w / h) - (width / height)) < 0.01
|
|
notes.append("%s: %dx%d -> %dx%d, %s%s"
|
|
% (label, w, h, width, height, behaviour,
|
|
"" if same_aspect else " — ASPECT DIFFERS, so this is visible"))
|
|
|
|
if first_frame is None and last_frame is None:
|
|
notes.append("no frames given: this is a text-to-video run, which is "
|
|
"legal here — an unrelated-looking result is not a bug")
|
|
if first_frame is not None:
|
|
describe(first_frame, "first_frame", "stretched to fit (core anchors geometry on it)")
|
|
if last_frame is not None:
|
|
describe(last_frame, "last_frame", "centre-cropped (core preserves its aspect)")
|
|
|
|
if first_frame is not None and last_frame is not None:
|
|
try:
|
|
a = (int(first_frame.shape[2]), int(first_frame.shape[1]))
|
|
b = (int(last_frame.shape[2]), int(last_frame.shape[1]))
|
|
except Exception:
|
|
a = b = None
|
|
if a and b and a != b:
|
|
notes.append("the two frames are different sizes (%dx%d and %dx%d) "
|
|
"— legal, but they will be fitted differently, so the "
|
|
"clip may drift in framing" % (a[0], a[1], b[0], b[1]))
|
|
return notes
|
|
|
|
|
|
# ── conditioning, sampling, decoding ─────────────────────────────────────────
|
|
def encode_fl2va(clip, vae, prompt, width, height, frames,
|
|
first_frame=None, last_frame=None):
|
|
"""Conditioning + empty AV latent, from core's own fl2va node.
|
|
|
|
Called rather than reimplemented: the keyframe encoding, the resize policy
|
|
and the latent geometry are the model author's code, and this pack has no
|
|
business having a second opinion about them.
|
|
"""
|
|
h3 = core_h3()
|
|
node = h3.MiniMaxH3ImageToVideo
|
|
fn = getattr(node, "execute", None) or getattr(node, "encode", None)
|
|
if fn is None:
|
|
_fail("core's MiniMax-H3 first/last-frame node changed shape in this "
|
|
"ComfyUI version. Update this pack (Manager > Update), or roll "
|
|
"ComfyUI back to the version the workflow page names.")
|
|
out = fn(clip=clip, vae=vae, prompt=prompt, width=width, height=height,
|
|
length=frames, first_frame=first_frame, last_frame=last_frame)
|
|
values = unwrap(out)
|
|
if len(values) < 2:
|
|
_fail("core's fl2va node returned %d value(s) where 2 were expected."
|
|
% len(values))
|
|
return values[0], values[1]
|
|
|
|
|
|
def sample_av(model, conditioning, latent, steps, seed):
|
|
"""Mirror of core's SamplerCustomAdvanced with a Basic guider.
|
|
|
|
Basic (not CFG) guidance is what the model's own template uses: this path
|
|
has no negative prompt, so there is no second branch to weigh against.
|
|
"""
|
|
import comfy.model_management
|
|
import comfy.sample
|
|
import comfy.samplers
|
|
import comfy.utils
|
|
import latent_preview
|
|
try:
|
|
from comfy_extras.nodes_custom_sampler import Guider_Basic, Noise_RandomNoise
|
|
except Exception:
|
|
_fail("this ComfyUI is missing the custom sampler internals this pack "
|
|
"drives. Update ComfyUI and restart.")
|
|
|
|
guider = Guider_Basic(model)
|
|
guider.set_conds(conditioning)
|
|
sampler = comfy.samplers.sampler_object(SAMPLER_NAME)
|
|
sigmas = comfy.samplers.calculate_sigmas(
|
|
model.get_model_object("model_sampling"), SCHEDULER, int(steps)).cpu()
|
|
noise = Noise_RandomNoise(int(seed))
|
|
|
|
latent = dict(latent)
|
|
samples_in = latent["samples"]
|
|
try:
|
|
samples_in = comfy.sample.fix_empty_latent_channels(
|
|
guider.model_patcher, samples_in,
|
|
latent.get("downscale_ratio_spacial", None),
|
|
latent.get("downscale_ratio_temporal", None))
|
|
except TypeError: # older signature took two arguments
|
|
samples_in = comfy.sample.fix_empty_latent_channels(
|
|
guider.model_patcher, samples_in)
|
|
latent["samples"] = samples_in
|
|
|
|
callback = latent_preview.prepare_callback(
|
|
guider.model_patcher, sigmas.shape[-1] - 1)
|
|
samples = guider.sample(
|
|
noise.generate_noise(latent), samples_in, sampler, sigmas,
|
|
denoise_mask=latent.get("noise_mask", None), callback=callback,
|
|
disable_pbar=not comfy.utils.PROGRESS_BAR_ENABLED, seed=noise.seed)
|
|
samples = samples.to(comfy.model_management.intermediate_device())
|
|
|
|
out = dict(latent)
|
|
out.pop("downscale_ratio_spacial", None)
|
|
out.pop("downscale_ratio_temporal", None)
|
|
out["samples"] = samples
|
|
return out
|
|
|
|
|
|
def decode_video(vae, latent):
|
|
"""Mirror of core's VAEDecode, including the nested-latent unbind.
|
|
|
|
An H3 latent is always a PAIR — video and audio in one NestedTensor — so
|
|
the video VAE gets element 0 and the audio VAE gets the last one.
|
|
"""
|
|
lat = latent["samples"]
|
|
if getattr(lat, "is_nested", False):
|
|
lat = lat.unbind()[0]
|
|
images = vae.decode(lat)
|
|
if len(images.shape) == 5:
|
|
images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1])
|
|
return images
|
|
|
|
|
|
def decode_audio(audio_vae, latent):
|
|
"""core's own `vae_decode_audio` — it normalises level, we should not."""
|
|
try:
|
|
from comfy_extras.nodes_audio import vae_decode_audio
|
|
except Exception:
|
|
_fail("this ComfyUI cannot decode audio latents (comfy_extras/"
|
|
"nodes_audio.py is missing or older than this pack expects). "
|
|
"Update ComfyUI and restart.")
|
|
return vae_decode_audio(audio_vae, latent)
|
|
|
|
|
|
def report(title, lines):
|
|
"""One block, printed and returned, so the run explains itself either way."""
|
|
body = "\n".join(" " + str(line) for line in lines if line)
|
|
text = "[SimpleShot] %s\n%s" % (title, body)
|
|
print(text, flush=True)
|
|
return text
|