Jason Hoku
55ac08a73b
Extends the May 2026 Florence2 fix (29GB -> <5GB) to every other out-of-graph
inference entry point in the project. Same root cause everywhere: we call
ComfyUI models from background threads / our orchestrator, bypassing
ComfyUI/execution.py:732 which wraps every node in torch.inference_mode().
Decorated with @torch.inference_mode() in image_generation.py:
- generate_image() — SDXL/diffusion sampling (nodes.common_ksampler)
- upscale_image() — KSampler hires fix + ImageUpscaleWithModel
(ESRGAN-style) + VAE encode/decode
- decode_latent_with_vae() — VAE decode (with NaN float32 retry)
- seedvr2_upscale() — SeedVR2 diffusion-based video upscaler
(DiT + VAE + sampling loop)
Plus prominent module-docstring warning blocks in:
- image_generation.py: full "why this matters / where / what to do if
you add a new entry point" explanation
- ltx_video_generation.py: warning that _call_node()'s inference_mode wrap
is load-bearing and must not be removed
Inline comment above each decorated function points back to the module
docstring so anyone editing one function still sees the rationale.
Updated the stale comment in decode_latent_with_vae explaining .detach()
(was 'tensor may have requires_grad=True' — under inference_mode it can't,
so .detach() is now a defensive no-op; kept for double-safety).
Tests: 118/118 pass.
717 lines
27 KiB
Python
717 lines
27 KiB
Python
"""
|
||
LTX 2.3 Video Generation Module
|
||
Two-stage SamplerCustomAdvanced pipeline with parallel audio rail.
|
||
|
||
============================================================================
|
||
🚨 CRITICAL: _call_node() WRAPS EVERY INVOCATION IN torch.inference_mode() 🚨
|
||
============================================================================
|
||
The _call_node() helper in THIS file is what every ComfyUI node call in USCG
|
||
goes through (florence2_hires.py and this file's 21 LTX node invocations).
|
||
That helper applies `torch.inference_mode()` around the underlying execute()
|
||
or FUNCTION call — mirroring ComfyUI/execution.py:732 which wraps every node
|
||
in the prompt executor's inference_mode block.
|
||
|
||
DO NOT remove the inference_mode wrap from _call_node. Without it the
|
||
SamplerCustomAdvanced stage 1+2, VAEDecodeTiled, LTXVLatentUpsampler, CLIP
|
||
encode calls, and all other model-forward nodes leak intermediate
|
||
activations across diffusion steps. We measured a 6× VRAM blowup on
|
||
Florence2 from this exact missing wrapper (29GB on a 16GB card vs <5GB
|
||
in the standalone workflow on the same image, May 2026). Same root cause
|
||
applies to every other model-inference path that doesn't go through
|
||
ComfyUI's prompt executor.
|
||
|
||
If you add new node invocations that DON'T go through _call_node (e.g.,
|
||
direct `.execute()` or `instance.FUNCTION()` calls), wrap them yourself
|
||
in `with torch.inference_mode():` — or better, route them through
|
||
_call_node so the wrapper is automatic.
|
||
============================================================================
|
||
|
||
|
||
Pinned LTX node pack version: TBD — set during first smoke test.
|
||
Required nodes (looked up via nodes.NODE_CLASS_MAPPINGS):
|
||
- DiffusionModelLoaderKJ
|
||
- DualCLIPLoader (built-in)
|
||
- VAELoaderKJ
|
||
- LatentUpscaleModelLoader
|
||
- LTXVPreprocess, LTXVImgToVideoInplace, LTXVConditioning,
|
||
- LTXVCropGuides, LTXVConcatAVLatent, LTXVSeparateAVLatent,
|
||
- LTXVEmptyLatentAudio, LTXVAudioVAEDecode,
|
||
- EmptyLTXVLatentVideo, LTXVLatentUpsampler,
|
||
- ManualSigmas, KSamplerSelect, RandomNoise, CFGGuider,
|
||
- SamplerCustomAdvanced, VAEDecodeTiled, CreateVideo, SaveVideo
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import time
|
||
import builtins
|
||
import uuid
|
||
from typing import List
|
||
|
||
|
||
def safe_print(*args, **kwargs):
|
||
"""Windows-safe print mirroring config_builder_node.safe_print."""
|
||
try:
|
||
builtins.print(*args, **kwargs)
|
||
except (OSError, ValueError):
|
||
try:
|
||
msg = " ".join(str(a) for a in args) + kwargs.get("end", "\n")
|
||
sys.__stdout__.write(msg)
|
||
sys.__stdout__.flush()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
print = safe_print
|
||
|
||
|
||
def parse_sigmas(sigma_str):
|
||
"""Parse a comma-separated sigma string into a list of floats.
|
||
|
||
Args:
|
||
sigma_str: e.g. "0.85, 0.7250, 0.4219, 0.0"
|
||
|
||
Returns:
|
||
List of floats.
|
||
|
||
Raises:
|
||
ValueError: empty string, fewer than 2 values, or any token not a valid float.
|
||
"""
|
||
if not sigma_str or not sigma_str.strip():
|
||
raise ValueError("Sigma string is empty")
|
||
|
||
tokens = [t.strip() for t in sigma_str.split(",")]
|
||
# Reject empty tokens (catches trailing/leading/double commas)
|
||
for tok in tokens:
|
||
if not tok:
|
||
raise ValueError(
|
||
"Sigma string has an empty token (trailing/leading/double comma): " + repr(sigma_str)
|
||
)
|
||
if len(tokens) < 2:
|
||
raise ValueError(
|
||
"Sigma string must contain at least 2 comma-separated values, got: " + repr(sigma_str)
|
||
)
|
||
|
||
result = []
|
||
for tok in tokens:
|
||
try:
|
||
result.append(float(tok))
|
||
except ValueError:
|
||
raise ValueError("Sigma token " + repr(tok) + " is not a valid float (in " + repr(sigma_str) + ")")
|
||
return result
|
||
|
||
|
||
import shutil
|
||
|
||
|
||
# Required LTX node class names — looked up via nodes.NODE_CLASS_MAPPINGS
|
||
REQUIRED_LTX_NODE_NAMES = [
|
||
"DiffusionModelLoaderKJ",
|
||
"DualCLIPLoader",
|
||
"VAELoaderKJ",
|
||
"LatentUpscaleModelLoader",
|
||
"LTXVPreprocess",
|
||
"LTXVImgToVideoInplace",
|
||
"LTXVConditioning",
|
||
"LTXVCropGuides",
|
||
"LTXVConcatAVLatent",
|
||
"LTXVSeparateAVLatent",
|
||
"LTXVEmptyLatentAudio",
|
||
"LTXVAudioVAEDecode",
|
||
"EmptyLTXVLatentVideo",
|
||
"LTXVLatentUpsampler",
|
||
"ManualSigmas",
|
||
"KSamplerSelect",
|
||
"RandomNoise",
|
||
"CFGGuider",
|
||
"SamplerCustomAdvanced",
|
||
"VAEDecodeTiled",
|
||
"CreateVideo",
|
||
"SaveVideo",
|
||
]
|
||
|
||
|
||
def get_ltx_node_classes():
|
||
"""Look up all required LTX nodes in NODE_CLASS_MAPPINGS.
|
||
|
||
Returns:
|
||
Dict mapping node name to class.
|
||
|
||
Raises:
|
||
RuntimeError: any required node is missing.
|
||
"""
|
||
import nodes
|
||
found = {}
|
||
missing = []
|
||
for name in REQUIRED_LTX_NODE_NAMES:
|
||
cls = nodes.NODE_CLASS_MAPPINGS.get(name)
|
||
if cls is None:
|
||
missing.append(name)
|
||
else:
|
||
found[name] = cls
|
||
if missing:
|
||
raise RuntimeError(
|
||
"LTX 2.3 video generation requires the following ComfyUI nodes "
|
||
"(install via Comfy Manager - search 'LTXVideo' and 'KJNodes'):\n"
|
||
+ "\n".join(" - " + n for n in missing)
|
||
)
|
||
return found
|
||
|
||
|
||
def _call_node(node_cls, **kwargs):
|
||
"""Call a Comfy node regardless of V1/V3 API. Returns whatever the node returns
|
||
(typically a tuple, but V3 nodes return an object with .output).
|
||
|
||
CRITICAL: Wraps the invocation in torch.inference_mode() to mirror ComfyUI's
|
||
PromptExecutor (ComfyUI/execution.py:732). ComfyUI's executor wraps every node
|
||
call in a global inference_mode block; when we call nodes directly via this
|
||
helper we bypass that wrapper, which leaves autograd's version-counter and
|
||
reference-keeping machinery active. For pure tensor-op nodes (GrowMask,
|
||
FeatherMask) that's a small overhead; for inference-heavy nodes (Florence2Run
|
||
beam search, LTX SamplerCustomAdvanced stage 1/2, VAEDecodeTiled, etc.) it
|
||
means intermediate activations and past_key_values can't be released between
|
||
steps — measured as a 6x VRAM blowup on Florence2 (29GB on a 16GB card vs <5GB
|
||
in the standalone workflow on the same image, May 2026). Wrapping here fixes
|
||
every existing and future _call_node site at once.
|
||
|
||
V3 (io.ComfyNode subclass): class-level .execute() classmethod — must be inside
|
||
a CurrentNodeContext block.
|
||
V1 (legacy): node has a FUNCTION attribute naming an instance method; instantiate
|
||
the class then call that method on the instance. KJ nodes and most core ComfyUI
|
||
nodes work this way.
|
||
"""
|
||
import torch as _torch_ic
|
||
with _torch_ic.inference_mode():
|
||
# Try V3 first via subclass check (avoids triggering AttributeError noise on V1 nodes
|
||
# that may also happen to have an unrelated `execute` attribute somewhere).
|
||
try:
|
||
from comfy_api.latest import io
|
||
if isinstance(node_cls, type) and issubclass(node_cls, io.ComfyNode):
|
||
return node_cls.execute(**kwargs)
|
||
except (ImportError, TypeError):
|
||
pass
|
||
|
||
# V1 path: FUNCTION attribute names the instance method to invoke
|
||
fn_name = getattr(node_cls, "FUNCTION", None)
|
||
if fn_name:
|
||
inst = node_cls()
|
||
return getattr(inst, fn_name)(**kwargs)
|
||
|
||
# Last-resort fallback: try class-level .execute()
|
||
if hasattr(node_cls, "execute"):
|
||
return node_cls.execute(**kwargs)
|
||
|
||
raise RuntimeError("Cannot determine entry method for node " + repr(node_cls))
|
||
|
||
|
||
def _unwrap(result, idx=0):
|
||
"""Normalize a node's return value to a single positional output by index.
|
||
|
||
V3 io.NodeOutput stores positional outputs in `.args` (a tuple) and supports
|
||
`result[idx]` via __getitem__. V1 nodes return plain tuples/lists. A few helpers
|
||
return a single object directly.
|
||
"""
|
||
# V3 io.NodeOutput
|
||
if hasattr(result, "args") and isinstance(getattr(result, "args"), tuple):
|
||
return result.args[idx]
|
||
# Some legacy helpers used a `.output` attribute (e.g. SeedVR2-style returns)
|
||
if hasattr(result, "output"):
|
||
out = result.output
|
||
if isinstance(out, (tuple, list)):
|
||
return out[idx]
|
||
return out
|
||
# V1 plain tuple/list
|
||
if isinstance(result, (tuple, list)):
|
||
return result[idx]
|
||
# Single bare value (unusual but possible)
|
||
return result
|
||
|
||
|
||
def preflight_ltx(config):
|
||
"""Validate that LTX gen can run for this config. Raises RuntimeError on any
|
||
missing node, missing model file, or missing ffmpeg.
|
||
|
||
Call ONCE per LTX grid run before any gen, and once more per-config for
|
||
things that vary per-config (i2v image existence, sigma string parses).
|
||
"""
|
||
import folder_paths
|
||
|
||
# 1. LTX node pack installed?
|
||
get_ltx_node_classes()
|
||
|
||
# 2. ffmpeg available?
|
||
if shutil.which("ffmpeg") is None:
|
||
raise RuntimeError(
|
||
"ffmpeg required for LTX video output. Install ffmpeg and ensure "
|
||
"it's on your system PATH."
|
||
)
|
||
|
||
# 3. Model files exist?
|
||
# Note: latent_upscaler lives in `latent_upscale_models/` (consumed by
|
||
# LatentUpscaleModelLoader), NOT `upscale_models/` (which is for image upscalers).
|
||
checks = [
|
||
("diffusion_models", config["model"]),
|
||
("text_encoders", config["clip_models"][0]),
|
||
("text_encoders", config["clip_models"][1]),
|
||
("vae", config["vae_video"]),
|
||
("vae", config["vae_audio"]),
|
||
("latent_upscale_models", config["latent_upscaler"]),
|
||
]
|
||
missing_files = []
|
||
for folder_key, name in checks:
|
||
path = folder_paths.get_full_path(folder_key, name)
|
||
if path is None:
|
||
missing_files.append(folder_key + "/" + name)
|
||
if missing_files:
|
||
raise RuntimeError(
|
||
"LTX model files not found:\n" + "\n".join(" - " + m for m in missing_files)
|
||
)
|
||
|
||
# 4. Sigma strings parse?
|
||
parse_sigmas(config["sigmas_stage1"])
|
||
parse_sigmas(config["sigmas_stage2"])
|
||
|
||
# 5. i2v image exists if set?
|
||
img = config.get("input_image")
|
||
if img:
|
||
# Resolve via ComfyUI's input-dir lookup if it's not an absolute path
|
||
# (drag-drop and clipboard-paste uploads land there as bare filenames).
|
||
resolved = img
|
||
if not os.path.isabs(img) and not os.path.isfile(img):
|
||
try:
|
||
import folder_paths as _fp
|
||
candidate = _fp.get_annotated_filepath(img)
|
||
if candidate and os.path.isfile(candidate):
|
||
resolved = candidate
|
||
except Exception:
|
||
pass
|
||
if not os.path.isfile(resolved):
|
||
raise RuntimeError("Input image not found: " + img)
|
||
|
||
# 6. Even dimensions (LTX divides by 2 for latent space)?
|
||
w = int(config["width"])
|
||
h = int(config["height"])
|
||
if w % 2 != 0 or h % 2 != 0:
|
||
raise RuntimeError("LTX requires even width/height, got " + str(w) + "x" + str(h))
|
||
|
||
|
||
def load_ltx_models(config):
|
||
"""Load all 5 LTX model files. Cached individually by file path.
|
||
|
||
Returns:
|
||
Dict with keys: diffusion_model, dual_clip, video_vae, audio_vae,
|
||
latent_upscaler.
|
||
|
||
Raises:
|
||
RuntimeError: on missing nodes (caught earlier in preflight) or load failure.
|
||
"""
|
||
from comfy_execution.utils import CurrentNodeContext
|
||
from . import model_cache
|
||
|
||
nodes_map = get_ltx_node_classes()
|
||
|
||
weight_dtype = config.get("weight_dtype", "default")
|
||
compute_dtype = config.get("compute_dtype", "default")
|
||
device = config.get("device", "main_device")
|
||
|
||
# Cache keys
|
||
diff_key = config["model"] + "::" + weight_dtype + "::" + compute_dtype
|
||
clip_key = config["clip_models"][0] + "::" + config["clip_models"][1] + "::ltxv::default"
|
||
vvae_key = config["vae_video"] + "::" + device + "::bf16"
|
||
avae_key = config["vae_audio"] + "::" + device + "::bf16"
|
||
upsc_key = config["latent_upscaler"]
|
||
|
||
dummy_prompt_id = str(uuid.uuid4())
|
||
|
||
out = {}
|
||
|
||
# --- Diffusion model (V1 API: DiffusionModelLoaderKJ.patch_and_load) ---
|
||
if diff_key in model_cache.ltx_diffusion_model_cache:
|
||
out["diffusion_model"] = model_cache.ltx_diffusion_model_cache[diff_key]
|
||
else:
|
||
with CurrentNodeContext(prompt_id=dummy_prompt_id, node_id="uscg_ltx_diff", list_index=0):
|
||
r = _call_node(
|
||
nodes_map["DiffusionModelLoaderKJ"],
|
||
model_name=config["model"],
|
||
weight_dtype=weight_dtype,
|
||
compute_dtype=compute_dtype,
|
||
patch_cublaslinear=False,
|
||
sage_attention="disabled",
|
||
enable_fp16_accumulation=False,
|
||
)
|
||
m = _unwrap(r, 0)
|
||
model_cache.ltx_diffusion_model_cache[diff_key] = m
|
||
model_cache._evict_to_max(model_cache.ltx_diffusion_model_cache, 1)
|
||
out["diffusion_model"] = m
|
||
|
||
# --- Dual CLIP (V1 API: DualCLIPLoader().load_clip) ---
|
||
if clip_key in model_cache.ltx_dual_clip_cache:
|
||
out["dual_clip"] = model_cache.ltx_dual_clip_cache[clip_key]
|
||
else:
|
||
r = _call_node(
|
||
nodes_map["DualCLIPLoader"],
|
||
clip_name1=config["clip_models"][0],
|
||
clip_name2=config["clip_models"][1],
|
||
type="ltxv",
|
||
device="default",
|
||
)
|
||
c = _unwrap(r, 0)
|
||
model_cache.ltx_dual_clip_cache[clip_key] = c
|
||
model_cache._evict_to_max(model_cache.ltx_dual_clip_cache, 2)
|
||
out["dual_clip"] = c
|
||
|
||
# --- Video VAE (V1 API: VAELoaderKJ.load_vae) ---
|
||
if vvae_key in model_cache.ltx_video_vae_cache:
|
||
out["video_vae"] = model_cache.ltx_video_vae_cache[vvae_key]
|
||
else:
|
||
with CurrentNodeContext(prompt_id=dummy_prompt_id, node_id="uscg_ltx_vvae", list_index=0):
|
||
r = _call_node(
|
||
nodes_map["VAELoaderKJ"],
|
||
vae_name=config["vae_video"],
|
||
device=device,
|
||
weight_dtype="bf16",
|
||
)
|
||
v = _unwrap(r, 0)
|
||
model_cache.ltx_video_vae_cache[vvae_key] = v
|
||
model_cache._evict_to_max(model_cache.ltx_video_vae_cache, 1)
|
||
out["video_vae"] = v
|
||
|
||
# --- Audio VAE (V1 API: VAELoaderKJ.load_vae) ---
|
||
if avae_key in model_cache.ltx_audio_vae_cache:
|
||
out["audio_vae"] = model_cache.ltx_audio_vae_cache[avae_key]
|
||
else:
|
||
with CurrentNodeContext(prompt_id=dummy_prompt_id, node_id="uscg_ltx_avae", list_index=0):
|
||
r = _call_node(
|
||
nodes_map["VAELoaderKJ"],
|
||
vae_name=config["vae_audio"],
|
||
device=device,
|
||
weight_dtype="bf16",
|
||
)
|
||
v = _unwrap(r, 0)
|
||
model_cache.ltx_audio_vae_cache[avae_key] = v
|
||
model_cache._evict_to_max(model_cache.ltx_audio_vae_cache, 1)
|
||
out["audio_vae"] = v
|
||
|
||
# --- Latent upscaler (V1 API: LatentUpscaleModelLoader().load_model) ---
|
||
if upsc_key in model_cache.ltx_latent_upscaler_cache:
|
||
out["latent_upscaler"] = model_cache.ltx_latent_upscaler_cache[upsc_key]
|
||
else:
|
||
r = _call_node(nodes_map["LatentUpscaleModelLoader"], model_name=config["latent_upscaler"])
|
||
u = _unwrap(r, 0)
|
||
model_cache.ltx_latent_upscaler_cache[upsc_key] = u
|
||
model_cache._evict_to_max(model_cache.ltx_latent_upscaler_cache, 1)
|
||
out["latent_upscaler"] = u
|
||
|
||
return out
|
||
|
||
|
||
def encode_ltx_prompts(dual_clip, positive_text, negative_text, frame_rate):
|
||
"""Dual-CLIP encode positive and negative prompts, then wrap with LTXVConditioning.
|
||
|
||
Returns:
|
||
Tuple (cond_pos, cond_neg) - LTXVConditioning-wrapped pair ready for
|
||
CFGGuider or LTXVCropGuides.
|
||
"""
|
||
nodes_map = get_ltx_node_classes()
|
||
|
||
# Standard CLIPTextEncode (built-in V1)
|
||
encoder = nodes_map.get("CLIPTextEncode")
|
||
if encoder is None:
|
||
import nodes
|
||
encoder = nodes.NODE_CLASS_MAPPINGS["CLIPTextEncode"]
|
||
|
||
pos = _unwrap(_call_node(encoder, clip=dual_clip, text=positive_text), 0)
|
||
neg = _unwrap(_call_node(encoder, clip=dual_clip, text=negative_text), 0)
|
||
|
||
# LTXVConditioning wraps with frame_rate (V3 API: ComfyNode subclass)
|
||
from comfy_execution.utils import CurrentNodeContext
|
||
with CurrentNodeContext(prompt_id=str(uuid.uuid4()), node_id="uscg_ltx_cond", list_index=0):
|
||
r = _call_node(nodes_map["LTXVConditioning"], positive=pos, negative=neg, frame_rate=int(frame_rate))
|
||
|
||
return _unwrap(r, 0), _unwrap(r, 1)
|
||
|
||
|
||
def ltx_video_generate(config, ltx_models, output_path):
|
||
"""Run the two-stage LTX pipeline and write an mp4.
|
||
|
||
Phase A: text-to-video only. `input_image` is ignored.
|
||
`audio_mode` is always treated as "on".
|
||
|
||
Args:
|
||
config: Per-config dict (single config from cartesian expansion).
|
||
ltx_models: Output of load_ltx_models().
|
||
output_path: Absolute path for the mp4 output (no extension required;
|
||
".mp4" appended if missing).
|
||
|
||
Returns:
|
||
Dict with manifest fields: video_path, frames, fps, duration, etc.
|
||
|
||
Raises:
|
||
Bubbles RuntimeError / OOM exceptions to caller.
|
||
"""
|
||
from comfy_execution.utils import CurrentNodeContext
|
||
import torch
|
||
|
||
t0 = time.time()
|
||
nodes_map = get_ltx_node_classes()
|
||
prompt_id = str(uuid.uuid4())
|
||
|
||
width = int(config["width"])
|
||
height = int(config["height"])
|
||
duration_seconds = int(config["duration_seconds"])
|
||
frame_rate = int(config["frame_rate"])
|
||
cfg_scale = float(config.get("cfg", 1.0))
|
||
seed_stage1 = int(config["seed"])
|
||
seed_stage2 = seed_stage1 + 1
|
||
|
||
sigmas1 = config["sigmas_stage1"]
|
||
sigmas2 = config["sigmas_stage2"]
|
||
sampler_stage1 = config.get("sampler_stage1", "euler_ancestral_cfg_pp")
|
||
sampler_stage2 = config.get("sampler_stage2", "euler_cfg_pp")
|
||
|
||
frames = duration_seconds * frame_rate + 1
|
||
|
||
print("[GridTester] LTX gen: " + str(width) + "x" + str(height) +
|
||
", dur=" + str(duration_seconds) + "s, fps=" + str(frame_rate) +
|
||
", frames=" + str(frames) + ", seed=" + str(seed_stage1))
|
||
|
||
# Encode prompts
|
||
cond_pos, cond_neg = encode_ltx_prompts(
|
||
ltx_models["dual_clip"],
|
||
config.get("positive", ""),
|
||
config.get("negative", ""),
|
||
frame_rate,
|
||
)
|
||
|
||
diff_model = ltx_models["diffusion_model"]
|
||
video_vae = ltx_models["video_vae"]
|
||
audio_vae = ltx_models["audio_vae"]
|
||
upscaler = ltx_models["latent_upscaler"]
|
||
|
||
# Run all node executions inside a single mock V3 context
|
||
with CurrentNodeContext(prompt_id=prompt_id, node_id="uscg_ltx_gen", list_index=0):
|
||
# 1. Empty video latent (dimensions /2 for latent space)
|
||
empty_video_latent = _unwrap(_call_node(
|
||
nodes_map["EmptyLTXVLatentVideo"],
|
||
width=width // 2, height=height // 2, length=frames, batch_size=1,
|
||
), 0)
|
||
|
||
# 2. Empty audio latent
|
||
empty_audio_latent = _unwrap(_call_node(
|
||
nodes_map["LTXVEmptyLatentAudio"],
|
||
frames_number=frames, frame_rate=frame_rate, batch_size=1, audio_vae=audio_vae,
|
||
), 0)
|
||
|
||
# 2.5. If i2v: preprocess input image and run LTXVImgToVideoInplace stage 1.
|
||
# The preprocessed image tensor is reused for stage 2 (post-upscale refinement).
|
||
input_image_path = config.get("input_image")
|
||
is_i2v = bool(input_image_path)
|
||
i2v_preprocessed_image = None
|
||
if is_i2v:
|
||
from PIL import Image as PILImage
|
||
import numpy as np
|
||
import torch as _torch
|
||
import os as _os
|
||
|
||
# Resolve the input_image. If it's an absolute path, use as-is. Otherwise
|
||
# try resolving via ComfyUI's standard input-dir lookup (drag-drop and
|
||
# clipboard-paste uploads land there as bare filenames).
|
||
resolved_image_path = input_image_path
|
||
if not _os.path.isabs(input_image_path) and not _os.path.isfile(input_image_path):
|
||
try:
|
||
import folder_paths
|
||
candidate = folder_paths.get_annotated_filepath(input_image_path)
|
||
if candidate and _os.path.isfile(candidate):
|
||
resolved_image_path = candidate
|
||
except Exception:
|
||
pass # fall back to original; PIL will raise if truly missing
|
||
|
||
pil = PILImage.open(resolved_image_path).convert("RGB")
|
||
# Resize to target dims (matches reference workflow's ResizeImageMaskNode)
|
||
pil = pil.resize((width, height), PILImage.LANCZOS)
|
||
img_np = np.array(pil).astype(np.float32) / 255.0
|
||
img_tensor = _torch.from_numpy(img_np).unsqueeze(0) # [1, H, W, 3]
|
||
|
||
i2v_preprocessed_image = _unwrap(_call_node(
|
||
nodes_map["LTXVPreprocess"],
|
||
image=img_tensor,
|
||
img_compression=int(config.get("img_compression", 18)),
|
||
), 0)
|
||
|
||
# Apply ImgToVideoInplace stage 1 (strength_stage1)
|
||
video_pre = _unwrap(_call_node(
|
||
nodes_map["LTXVImgToVideoInplace"],
|
||
strength=float(config.get("image_strength_stage1", 0.8)),
|
||
bypass=False,
|
||
vae=video_vae,
|
||
image=i2v_preprocessed_image,
|
||
latent=empty_video_latent,
|
||
), 0)
|
||
else:
|
||
video_pre = empty_video_latent
|
||
|
||
# 3. Concat AV (stage 1 input)
|
||
stage1_input_latent = _unwrap(_call_node(
|
||
nodes_map["LTXVConcatAVLatent"],
|
||
video_latent=video_pre, audio_latent=empty_audio_latent,
|
||
), 0)
|
||
|
||
# 4. Stage 1 sampling
|
||
sampler1_obj = _unwrap(_call_node(nodes_map["KSamplerSelect"], sampler_name=sampler_stage1), 0)
|
||
|
||
sigmas1_tensor = _unwrap(_call_node(nodes_map["ManualSigmas"], sigmas=sigmas1), 0)
|
||
|
||
noise1_obj = _unwrap(_call_node(nodes_map["RandomNoise"], noise_seed=seed_stage1), 0)
|
||
|
||
guider1_obj = _unwrap(
|
||
_call_node(nodes_map["CFGGuider"],
|
||
model=diff_model, positive=cond_pos, negative=cond_neg, cfg=cfg_scale),
|
||
0,
|
||
)
|
||
|
||
sampled1_latent = _unwrap(
|
||
_call_node(nodes_map["SamplerCustomAdvanced"],
|
||
noise=noise1_obj, guider=guider1_obj, sampler=sampler1_obj,
|
||
sigmas=sigmas1_tensor, latent_image=stage1_input_latent),
|
||
0,
|
||
)
|
||
|
||
# 5. Separate AV
|
||
sep1 = _call_node(nodes_map["LTXVSeparateAVLatent"], av_latent=sampled1_latent)
|
||
video1, audio1 = _unwrap(sep1, 0), _unwrap(sep1, 1)
|
||
|
||
# 6. Spatial upscale of video latent
|
||
upscaled_video = _unwrap(
|
||
_call_node(nodes_map["LTXVLatentUpsampler"],
|
||
samples=video1, upscale_model=upscaler, vae=video_vae),
|
||
0,
|
||
)
|
||
|
||
# 7. If i2v: apply LTXVImgToVideoInplace stage 2 to the upscaled latent for refinement
|
||
if is_i2v and i2v_preprocessed_image is not None:
|
||
upscaled_video = _unwrap(_call_node(
|
||
nodes_map["LTXVImgToVideoInplace"],
|
||
strength=float(config.get("image_strength_stage2", 1.0)),
|
||
bypass=False,
|
||
vae=video_vae,
|
||
image=i2v_preprocessed_image,
|
||
latent=upscaled_video,
|
||
), 0)
|
||
|
||
# 8. Crop guides on the (possibly i2v-refined) upscaled latent
|
||
crop = _call_node(nodes_map["LTXVCropGuides"],
|
||
positive=cond_pos, negative=cond_neg, latent=upscaled_video)
|
||
crop_pos, crop_neg = _unwrap(crop, 0), _unwrap(crop, 1)
|
||
|
||
# 9. Reconcat AV for stage 2
|
||
stage2_input_latent = _unwrap(
|
||
_call_node(nodes_map["LTXVConcatAVLatent"],
|
||
video_latent=upscaled_video, audio_latent=audio1),
|
||
0,
|
||
)
|
||
|
||
# 10. Stage 2 sampling
|
||
sampler2_obj = _unwrap(_call_node(nodes_map["KSamplerSelect"], sampler_name=sampler_stage2), 0)
|
||
|
||
sigmas2_tensor = _unwrap(_call_node(nodes_map["ManualSigmas"], sigmas=sigmas2), 0)
|
||
|
||
noise2_obj = _unwrap(_call_node(nodes_map["RandomNoise"], noise_seed=seed_stage2), 0)
|
||
|
||
guider2_obj = _unwrap(
|
||
_call_node(nodes_map["CFGGuider"],
|
||
model=diff_model, positive=crop_pos, negative=crop_neg, cfg=cfg_scale),
|
||
0,
|
||
)
|
||
|
||
sampled2_latent = _unwrap(
|
||
_call_node(nodes_map["SamplerCustomAdvanced"],
|
||
noise=noise2_obj, guider=guider2_obj, sampler=sampler2_obj,
|
||
sigmas=sigmas2_tensor, latent_image=stage2_input_latent),
|
||
0,
|
||
)
|
||
|
||
# 11. Separate AV final
|
||
sep2 = _call_node(nodes_map["LTXVSeparateAVLatent"], av_latent=sampled2_latent)
|
||
video2, audio2 = _unwrap(sep2, 0), _unwrap(sep2, 1)
|
||
|
||
# 12. Decode video (tiled - non-tiled OOMs at typical durations)
|
||
frames_tensor = _unwrap(
|
||
_call_node(nodes_map["VAEDecodeTiled"],
|
||
samples=video2, vae=video_vae,
|
||
tile_size=768, overlap=64, temporal_size=4096, temporal_overlap=4),
|
||
0,
|
||
)
|
||
|
||
# 13. Decode audio (Phase A: always on)
|
||
audio_waveform = _unwrap(
|
||
_call_node(nodes_map["LTXVAudioVAEDecode"], samples=audio2, audio_vae=audio_vae),
|
||
0,
|
||
)
|
||
|
||
# 14. Create video container
|
||
video_obj = _unwrap(
|
||
_call_node(nodes_map["CreateVideo"], fps=frame_rate, images=frames_tensor, audio=audio_waveform),
|
||
0,
|
||
)
|
||
|
||
# 15. Save mp4 — call video.save_to() directly instead of going through
|
||
# the SaveVideo node. SaveVideo accesses cls.hidden.extra_pnginfo/prompt
|
||
# which the normal Comfy executor populates but we don't (and we don't
|
||
# need pnginfo metadata for video files anyway). save_to() is what
|
||
# SaveVideo.execute() calls internally — same code path, no hidden deps.
|
||
if not output_path.lower().endswith(".mp4"):
|
||
output_path = output_path + ".mp4"
|
||
out_dir = os.path.dirname(output_path)
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
|
||
try:
|
||
from comfy_api.latest._input_impl import VideoContainer as _VideoContainerEnum
|
||
except ImportError:
|
||
try:
|
||
from comfy_api.latest._input.video_types import VideoContainer as _VideoContainerEnum
|
||
except ImportError:
|
||
_VideoContainerEnum = None
|
||
|
||
save_kwargs = {"codec": "auto", "metadata": None}
|
||
if _VideoContainerEnum is not None:
|
||
try:
|
||
save_kwargs["format"] = _VideoContainerEnum("mp4")
|
||
except Exception:
|
||
save_kwargs["format"] = "mp4"
|
||
else:
|
||
save_kwargs["format"] = "mp4"
|
||
|
||
try:
|
||
video_obj.save_to(output_path, **save_kwargs)
|
||
except TypeError:
|
||
# Older Comfy: positional or differently-named kwargs
|
||
video_obj.save_to(output_path)
|
||
|
||
# Preview thumbnail: middle frame as PNG for fast dashboard previews
|
||
preview_path = None
|
||
try:
|
||
from PIL import Image as PILImage
|
||
import numpy as np
|
||
# frames_tensor shape: [N_frames, H, W, 3] float32 in [0, 1]
|
||
n = frames_tensor.shape[0]
|
||
mid_idx = n // 2
|
||
mid_np = (frames_tensor[mid_idx].cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
|
||
preview_path = output_path[:-4] + ".preview.png" if output_path.lower().endswith(".mp4") else output_path + ".preview.png"
|
||
PILImage.fromarray(mid_np, "RGB").save(preview_path)
|
||
except Exception as e:
|
||
print("[GridTester] Preview thumbnail write failed: " + str(e))
|
||
preview_path = None
|
||
|
||
duration = round(time.time() - t0, 2)
|
||
|
||
return {
|
||
"video_path": output_path,
|
||
"preview_path": preview_path,
|
||
"frames": frames,
|
||
"fps": frame_rate,
|
||
"duration_seconds": duration_seconds,
|
||
"width": width,
|
||
"height": height,
|
||
"duration": duration, # gen wall time
|
||
}
|